Strenghten preconditions of linalg.cross

This makes `linalg.cross` array API complaint (data-apis/array-api#415) and fixes a few bugs.

Fixes #77629
Fixes #83756

[ghstack-poisoned]

Author: lezcano

aten/src/ATen/TensorMeta.h

@@ -71,6 +71,7 @@ namespace impl {
 struct TORCH_API MetaBase {
   virtual const Tensor& maybe_get_output(int64_t output_idx) = 0;
 
+  // Note: [set_output_*]
   // See: https://github.com/pytorch/pytorch/issues/69813
   // Whenever defining the output properties in the META function of a
   // structured kernel (what was usually done with `set_output`), use one of

aten/src/ATen/native/Cross.cpp

@@ -9,17 +9,22 @@
 namespace at {
 namespace meta {
 
-TORCH_PRECOMPUTE_META_FUNC(linalg_cross)
-(const Tensor & input, const Tensor & other, const int64_t dimension) {
-  auto out_size = infer_size(input.sizes(), other.sizes());
-  Tensor input_broadcasted = input.expand(out_size);
-  Tensor other_broadcasted = other.expand(out_size);
+TORCH_META_FUNC(linalg_cross)
+(const Tensor & input, const Tensor & other, int64_t dim) {
+  auto x_d = input.dim();
+  auto y_d = other.dim();
+  // This is to avoid things like:
+  // linalg.cross(torch.randn(2, 3), torch.randn(5, 2, 3), dim=2)
+  // This would be very odd. This will still be possible by doing
+  // linalg.cross(torch.randn(2, 3).unsqueeze(0), torch.randn(5, 2, 3), dim=2)
+  TORCH_CHECK(x_d == y_d, "linalg.cross: inputs must have the same number of dimensions. Got ");
+  TORCH_CHECK(input.size(dim) == 3 && other.size(dim) == 3, "linalg.cross: inputs dimension ", dim, " must have length 3. Got ", input.size(dim), " and ", other.size(dim));
 
-  int64_t dim = maybe_wrap_dim(dimension, input.dim()); // default dim = -1
-  TORCH_CHECK(input_broadcasted.size(dim) == 3, "dimension ", dimension, " does not have size 3");
+  // Broadcast the batch dimension of input and other.
+  // Since the non-batch dimensions agree, this is the same as broadcast all the inputs
+  auto out_size = infer_size(input.sizes(), other.sizes());
 
   set_output_raw_strided(0, out_size, {}, input.options());
-  return TORCH_PRECOMPUTE_STRUCT(linalg_cross)().set_dim(dim);
 }
 
 }
@@ -56,8 +61,9 @@ Tensor & cross_out(const Tensor & input, const Tensor & other, const c10::option
 
 
 TORCH_IMPL_FUNC(linalg_cross_out)
-(const Tensor & input, const Tensor & other, const int64_t dim, const Tensor & out) {
-  auto out_size = infer_size(input.sizes(), other.sizes());
+(const Tensor & input, const Tensor & other, int64_t dim, const Tensor & out) {
+  dim = maybe_wrap_dim(dim, input.dim());
+  auto out_size = out.sizes();
   Tensor input_broadcasted = input.expand(out_size);
   Tensor other_broadcasted = other.expand(out_size);
 

aten/src/ATen/native/native_functions.yaml

@@ -12177,8 +12177,6 @@
 - func: linalg_cross.out(Tensor self, Tensor other, *, int dim=-1, Tensor(a!) out) -> Tensor(a!)
   python_module: linalg
   structured: True
-  precomputed:
-  - dim -> int dim
   dispatch:
     CPU, CUDA: linalg_cross_out
 

torch/testing/_internal/common_methods_invocations.py

@@ -109,6 +109,7 @@
     sample_inputs_linalg_cholesky,
     sample_inputs_linalg_cholesky_inverse,
     sample_inputs_cross,
+    error_inputs_cross,
     sample_inputs_linalg_qr_geqrf,
     sample_inputs_linalg_invertible,
     sample_inputs_lu_solve,
@@ -4925,7 +4926,6 @@ def _clamp_numpy(a, min=None, max=None):
 
     return np.minimum(max, np.maximum(a, min))
 
-
 def sample_inputs_cumprod(op_info, device, dtype, requires_grad, **kwargs):
     def make_arg(shape):
         # shrink values to be in the interval [-1, +1] for better precision in gradgradcheck

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

@@ -118,33 +118,21 @@ def fn_UVh(usv):
 
 
 def sample_inputs_cross(op_info, device, dtype, requires_grad, **kwargs):
-    yield SampleInput(
-        make_tensor((S, 3), device=device, dtype=dtype, requires_grad=requires_grad),
-        args=(
-            make_tensor(
-                (S, 3), device=device, dtype=dtype, requires_grad=requires_grad
-            ),
-        ),
-    )
-    yield SampleInput(
-        make_tensor((S, 3, S), device=device, dtype=dtype, requires_grad=requires_grad),
-        args=(
-            make_tensor(
-                (S, 3, S), device=device, dtype=dtype, requires_grad=requires_grad
-            ),
-        ),
-        kwargs={"dim": 1},
-    )
-    yield SampleInput(
-        make_tensor((S, 3), device=device, dtype=dtype, requires_grad=requires_grad),
-        args=(
-            make_tensor(
-                (S, 3), device=device, dtype=dtype, requires_grad=requires_grad
-            ),
-        ),
-        kwargs={"dim": -1},
-    )
+    make_arg = partial(make_tensor, dtype=dtype, device=device, requires_grad=requires_grad)
+    yield SampleInput(make_arg((S, 3)), args=(make_arg((S, 3)),))
+    yield SampleInput(make_arg((S, 3, S)), args=(make_arg((S, 3, S)),), kwargs=dict(dim=1))
+    yield SampleInput(make_arg((1, 3)), args=(make_arg((S, 3)),), kwargs=dict(dim=-1))
+
+def error_inputs_cross(op_info, device, dtype, requires_grad, **kwargs):
+    make_arg = partial(make_tensor, dtype=dtype, device=device, requires_grad=requires_grad)
+
+    sample = SampleInput(input=make_arg((S, 3)), args=(make_arg((S, 1)),))
+    err = "inputs dimension -1 must have length 3"
+    yield ErrorInput(sample, error_regex=err, error_type=RuntimeError)
 
+    sample = SampleInput(input=make_arg((5, S, 3)), args=(make_arg((S, 3)),))
+    err = "inputs must have the same number of dimensions"
+    yield ErrorInput(sample, error_regex=err, error_type=RuntimeError)
 
 def sample_inputs_householder_product(op_info, device, dtype, requires_grad, **kwargs):
     """