Patch pytorch for CVE-2025-55552

Author: archana25-ms

SPECS/pytorch/CVE-2025-55552-1.patch

@@ -0,0 +1,199 @@
+From 5b902cf87f2f8caa8a7023991992e547b1b7aef8 Mon Sep 17 00:00:00 2001
+From: Archana Shettigar <[email protected]>
+Date: Mon, 24 Nov 2025 11:08:31 +0530
+Subject: [PATCH] Addressing CVE-2025-55552
+
+---
+ test/inductor/test_torchinductor.py           | 22 ++++-
+ ...st_torchinductor_codegen_dynamic_shapes.py |  1 +
+ torch/_inductor/decomposition.py              | 95 +++++++++----------
+ 3 files changed, 66 insertions(+), 52 deletions(-)
+
+diff --git a/test/inductor/test_torchinductor.py b/test/inductor/test_torchinductor.py
+index f1cfb90c..4c424157 100644
+--- a/test/inductor/test_torchinductor.py
++++ b/test/inductor/test_torchinductor.py
+@@ -5872,7 +5872,7 @@ class CommonTemplate:
+         model = Model()
+         x = torch.rand(10, 3, 0)
+ 
+-        self.common(model, (x,))
++        self.common(model, (x,), exact_stride=True)
+ 
+     def test_randint(self):
+         @torch.compile(fullgraph=True)
+@@ -5907,9 +5907,21 @@ class CommonTemplate:
+     @config.patch(fallback_random=True)
+     def test_like_rands(self):
+         def fn(x):
+-            return torch.rand_like(x), torch.randn_like(x)
++            return torch.rand_like(x), torch.randn_like(x), torch.randint_like(x, 1, 11)
+ 
+-        self.common(fn, [torch.zeros([20, 20])])
++        self.common(fn, [torch.zeros([20, 20])], exact_stride=True)
++
++    @config.patch(fallback_random=True)
++    @xfail_if_mps  # 100% are not close
++    def test_like_rands_sliced(self):
++        def fn(x):
++            return (
++                torch.randn_like(x),
++                torch.randn_like(x),
++                torch.randint_like(x, 1, 11),
++            )
++
++        self.common(fn, (torch.zeros([3, 4])[:, ::2].permute(1, 0),), exact_stride=True)
+ 
+     def test_like_rands2(self):
+         # rand_like with kwargs `device` of str type
+@@ -5924,6 +5936,8 @@ class CommonTemplate:
+         a0 = fn(x).clone()
+         a1 = fn(x).clone()
+         self.assertFalse(torch.allclose(a0, a1))
++        self.assertEqual(a0.shape, a1.shape)
++        self.assertEqual(a0.stride(), a1.stride())
+ 
+     @requires_cuda()
+     def test_like_rands3(self):
+@@ -5940,6 +5954,8 @@ class CommonTemplate:
+         a1 = test_like_rands_on_different_device("cuda", "cpu")
+         self.assertTrue(a0.device.type == "cuda")
+         self.assertTrue(a1.device.type == "cpu")
++        self.assertEqual(a0.shape, a1.shape)
++        self.assertEqual(a0.stride(), a1.stride())
+ 
+     def test_max_pool2d_with_indices_backward(self):
+         def fn(a, b, c):
+diff --git a/test/inductor/test_torchinductor_codegen_dynamic_shapes.py b/test/inductor/test_torchinductor_codegen_dynamic_shapes.py
+index fa4b8040..ae52a802 100644
+--- a/test/inductor/test_torchinductor_codegen_dynamic_shapes.py
++++ b/test/inductor/test_torchinductor_codegen_dynamic_shapes.py
+@@ -162,6 +162,7 @@ test_failures = {
+     "test_bucketize_default_kwargs_dynamic_shapes": TestFailure("cpu"),
+     "test_bucketize_int_dynamic_shapes": TestFailure("cpu"),
+     "test_like_rands_dynamic_shapes": TestFailure(("cpu", "cuda")),
++    "test_like_rands_sliced_dynamic_shapes": TestFailure(("cpu", "cuda", "xpu")),
+     "test_linspace2_dynamic_shapes": TestFailure(("cpu", "cuda")),
+     "test_linspace3_dynamic_shapes": TestFailure(("cpu", "cuda")),
+     "test_max_pool2d6_dynamic_shapes": TestFailure(("cpu", "cuda")),
+diff --git a/torch/_inductor/decomposition.py b/torch/_inductor/decomposition.py
+index 88a56dea..c7943b2b 100644
+--- a/torch/_inductor/decomposition.py
++++ b/torch/_inductor/decomposition.py
+@@ -344,36 +344,6 @@ def view_copy_dtype(self, dtype):
+     return self.to(dtype).clone()
+ 
+ 
+-def get_like_layout(
+-    tensor: torch.Tensor, memory_format: Optional[torch.memory_format]
+-) -> torch.memory_format:
+-    # TODO: _to_copy tensor to stride permutation
+-    if memory_format is torch.preserve_format or memory_format is None:
+-        return utils.suggest_memory_format(tensor)
+-    else:
+-        return memory_format
+-
+-
+-@register_decomposition(aten.rand_like)
+-def rand_like(self, *, dtype=None, device=None, memory_format=None, **kwargs):
+-    return torch.rand(
+-        [*self.size()],
+-        dtype=dtype or self.dtype,
+-        device=device or self.device,
+-        **kwargs,
+-    ).to(memory_format=get_like_layout(self, memory_format))
+-
+-
+-@register_decomposition(aten.randn_like)
+-def randn_like(self, *, dtype=None, device=None, memory_format=None, **kwargs):
+-    return torch.randn(
+-        [*self.size()],
+-        dtype=dtype or self.dtype,
+-        device=device or self.device,
+-        **kwargs,
+-    ).to(memory_format=get_like_layout(self, memory_format))
+-
+-
+ @register_decomposition(aten.full_like)
+ def full_like(
+     self,
+@@ -396,30 +366,57 @@ def full_like(
+     ).to(memory_format=get_like_layout(self, memory_format))
+ 
+ 
+-@register_decomposition(aten.randint_like.default)
+-def randint_like(self, high, *, dtype=None, device=None, memory_format=None, **kwargs):
+-    return aten.randint.low(
+-        0,
+-        high,
+-        [*self.size()],
+-        dtype=dtype or self.dtype,
+-        device=device or self.device,
++def _rand_like(
++    rand_fn: Callable[..., torch.Tensor],
++    self: torch.Tensor,
++    *,
++    dtype: Optional[torch.dtype] = None,
++    device: Optional[torch.device] = None,
++    memory_format: torch.memory_format = torch.preserve_format,
++    **kwargs: Any,
++) -> torch.Tensor:
++    dtype = self.dtype if dtype is None else dtype
++    device = self.device if device is None else device
++
++    if memory_format != torch.preserve_format:
++        return rand_fn(
++            self.shape,
++            dtype=dtype,
++            device=device,
++            **kwargs,
++        ).to(memory_format=memory_format)
++
++    shape, permutation = _get_shape_permutation_like(self)
++    result = rand_fn(
++        shape,
++        dtype=dtype,
++        device=device,
+         **kwargs,
+-    ).to(memory_format=get_like_layout(self, memory_format))
++    )
++    if permutation == list(range(len(permutation))):
++        return result
++    return result.permute(permutation).clone()
+ 
+ 
++@register_decomposition(aten.rand_like)
++def rand_like(self: torch.Tensor, **kwargs: Any) -> torch.Tensor:
++    return _rand_like(torch.rand, self, **kwargs)
++
++
++@register_decomposition(aten.randn_like)
++def randn_like(self: torch.Tensor, **kwargs: Any) -> torch.Tensor:
++    return _rand_like(torch.randn, self, **kwargs)
++
++
++@register_decomposition(aten.randint_like.default)
++def randint_like(self: torch.Tensor, high: int, **kwargs: Any) -> torch.Tensor:
++    return _rand_like(functools.partial(aten.randint.low, 0, high), self, **kwargs)
++
+ @register_decomposition(aten.randint_like.low_dtype)
+ def randint_like_low(
+-    self, low, high, *, dtype=None, device=None, memory_format=None, **kwargs
+-):
+-    return aten.randint.low(
+-        low,
+-        high,
+-        [*self.size()],
+-        dtype=dtype or self.dtype,
+-        device=device or self.device,
+-        **kwargs,
+-    ).to(memory_format=get_like_layout(self, memory_format))
++    self: torch.Tensor, low: int, high: int, **kwargs: Any
++) -> torch.Tensor:
++    return _rand_like(functools.partial(aten.randint.low, low, high), self, **kwargs)
+ 
+ 
+ @register_decomposition(aten.randint.default)
+-- 
+2.45.4
+