OpInfo: use functools.partial to decrease noise in make_tensor calls

torch/testing/_internal/common_methods_invocations.py

@@ -1058,6 +1058,7 @@ def sample_inputs_baddbmm(op_info, device, dtype, requires_grad, **kwargs):
         ).with_metadata(broadcasts_input=broadcasts_input)
 
         if dtype.is_complex:
+            args = (make_arg(input_shape), make_arg(batch1_shape), make_arg(batch2_shape))
             yield SampleInput(
                 make_arg(input_shape),
                 make_arg(batch1_shape),

torch/testing/_internal/opinfo/definitions/_masked.py

@@ -56,35 +56,22 @@ def sample_inputs_softmax_variant(
 
 
 def _generate_masked_op_mask(input_shape, device, **kwargs):
+    make_arg = partial(
+        make_tensor, dtype=torch.bool, device=device, requires_grad=False
+    )
     yield None
-    yield make_tensor(input_shape, dtype=torch.bool, device=device, requires_grad=False)
+    yield make_arg(input_shape)
     if len(input_shape) > 2:
         # broadcast last mask dimension:
-        yield make_tensor(
-            input_shape[:-1] + (1,),
-            dtype=torch.bool,
-            device=device,
-            requires_grad=False,
-        )
+        yield make_arg(input_shape[:-1] + (1,))
         # broadcast middle mask dimension:
-        yield make_tensor(
-            input_shape[:1] + (1,) + input_shape[2:],
-            dtype=torch.bool,
-            device=device,
-            requires_grad=False,
-        )
+        yield make_arg(input_shape[:1] + (1,) + input_shape[2:])
         # broadcast first mask dimension:
-        yield make_tensor(
-            (1,) + input_shape[1:], dtype=torch.bool, device=device, requires_grad=False
-        )
+        yield make_arg((1,) + input_shape[1:])
         # mask.ndim < input.ndim
-        yield make_tensor(
-            input_shape[1:], dtype=torch.bool, device=device, requires_grad=False
-        )
+        yield make_arg(input_shape[1:])
         # mask.ndim == 1
-        yield make_tensor(
-            input_shape[-1:], dtype=torch.bool, device=device, requires_grad=False
-        )
+        yield make_arg(input_shape[-1:])
         # masks that require broadcasting of inputs (mask.ndim >
         # input.ndim) will not be supported, however, we may
         # reconsider this if there will be demand on this kind of
@@ -351,19 +338,16 @@ def sample_inputs_masked_logaddexp(op_info, device, dtype, requires_grad, **kwar
         list(_generate_masked_op_mask(shape, device, **kwargs)) for shape in shapes
     ]
 
+    make_arg = partial(
+        make_tensor, dtype=dtype, device=device, requires_grad=requires_grad
+    )
     for shape, input_masks, other_masks in zip(
         shapes, input_mask_lists, other_mask_lists
     ):
         for input_mask, other_mask in zip(input_masks, other_masks):
-            input = make_tensor(
-                shape, dtype=dtype, device=device, requires_grad=requires_grad
-            )
-            other = make_tensor(
-                shape, dtype=dtype, device=device, requires_grad=requires_grad
-            )
             yield SampleInput(
-                input.clone().requires_grad_(requires_grad),
-                args=(other.clone().requires_grad_(requires_grad),),
+                make_arg(shape),
+                args=(make_arg(shape),),
                 kwargs=dict(input_mask=input_mask, other_mask=other_mask),
             )
 

torch/testing/_internal/opinfo/definitions/linalg.py

@@ -374,6 +374,15 @@ def sample_inputs_linalg_norm(
     else:
         matrix_ords = (None, "fro", "nuc", inf, -inf, 1, -1, 2, -2)
 
+    make_arg = partial(
+        make_tensor,
+        dtype=dtype,
+        device=device,
+        requires_grad=requires_grad,
+        low=None,
+        high=None,
+    )
+
     for test_size in test_sizes:
         is_vector_norm = len(test_size) == 1
         is_matrix_norm = len(test_size) == 2
@@ -383,17 +392,7 @@ def sample_inputs_linalg_norm(
 
         for keepdim in [False, True]:
             if variant != "subgradient_at_zero" and is_valid_for_p2:
-                yield SampleInput(
-                    make_tensor(
-                        test_size,
-                        dtype=dtype,
-                        device=device,
-                        low=None,
-                        high=None,
-                        requires_grad=requires_grad,
-                    ),
-                    kwargs=dict(keepdim=keepdim),
-                )
+                yield SampleInput(make_arg(test_size), kwargs=dict(keepdim=keepdim))
 
             if not (is_vector_norm or is_matrix_norm):
                 continue
@@ -436,28 +435,12 @@ def sample_inputs_linalg_norm(
                     )
                 else:
                     yield SampleInput(
-                        make_tensor(
-                            test_size,
-                            dtype=dtype,
-                            device=device,
-                            low=None,
-                            high=None,
-                            requires_grad=requires_grad,
-                        ),
-                        args=(ord,),
-                        kwargs=dict(keepdim=keepdim),
+                        make_arg(test_size), args=(ord,), kwargs=dict(keepdim=keepdim)
                     )
 
                     if ord in ["nuc", "fro"]:
                         yield SampleInput(
-                            make_tensor(
-                                test_size,
-                                dtype=dtype,
-                                device=device,
-                                low=None,
-                                high=None,
-                                requires_grad=requires_grad,
-                            ),
+                            make_arg(test_size),
                             kwargs=dict(ord=ord, keepdim=keepdim, dim=(0, 1)),
                         )
 
@@ -732,16 +715,14 @@ def sample_inputs_linalg_ldl_solve(
     )
 
     # Symmetric case
+    make_arg = partial(
+        make_tensor, device=device, dtype=dtype, requires_grad=requires_grad
+    )
     for test_case in test_cases1:
         factors, pivots, _ = test_case
         factors.requires_grad = requires_grad
         for B_batch_shape in ((), factors.shape[:-2]):
-            B = make_tensor(
-                (*B_batch_shape, factors.shape[-1], S),
-                device=device,
-                dtype=dtype,
-                requires_grad=requires_grad,
-            )
+            B = make_arg((*B_batch_shape, factors.shape[-1], S))
             yield SampleInput(factors, args=(pivots, B), kwargs=dict(hermitian=False))
             clone_factors = factors.detach().clone().requires_grad_(requires_grad)
             yield SampleInput(
@@ -753,12 +734,7 @@ def sample_inputs_linalg_ldl_solve(
         factors, pivots, _ = test_case
         factors.requires_grad = requires_grad
         for B_batch_shape in ((), factors.shape[:-2]):
-            B = make_tensor(
-                (*B_batch_shape, factors.shape[-1], S),
-                device=device,
-                dtype=dtype,
-                requires_grad=requires_grad,
-            )
+            B = make_arg((*B_batch_shape, factors.shape[-1], S))
             yield SampleInput(factors, args=(pivots, B), kwargs=dict(hermitian=True))
             clone_factors = factors.detach().clone().requires_grad_(requires_grad)
             yield SampleInput(
@@ -1080,13 +1056,12 @@ def sample_inputs_tensorsolve(op_info, device, dtype, requires_grad, **kwargs):
     # a_shapes += [(0, 0, 1, 2, 3, 0)]
     dimss = [None, (0, 2)]
 
+    make_arg = partial(
+        make_tensor, dtype=dtype, device=device, requires_grad=requires_grad
+    )
     for a_shape, dims in itertools.product(a_shapes, dimss):
-        a = make_tensor(
-            a_shape, dtype=dtype, device=device, requires_grad=requires_grad
-        )
-        b = make_tensor(
-            a_shape[:2], dtype=dtype, device=device, requires_grad=requires_grad
-        )
+        a = make_arg(a_shape)
+        b = make_arg(a_shape[:2])
         yield SampleInput(a, args=(b,), kwargs=dict(dims=dims))