tensordot and vecdot array API conformity changes (#1577)

* Permit negative indices for tensordot when `axes` is a tuple

Per array API standard

Also fixes vecdot attempting to call `py_dot` with incorrect batching dims

* Fixes incorrect shape and stride in contiguous dot product indexers

* Fixes typo in GemmBatchNoAtomicFunctorThreadK

Closes #1570

* Aligns `vecdot` with array API spec changes

Only negative values for `axis` are permitted to avoid ambiguity

Now separately checks that the `axis` parameter is valid for each array before broadcasting occurs

* `test_usm_ndarray_linalg` changed to reflect `vecdot` and `tensordot` changes

* Updates `tensordot` and `vecdot` docstrings to reflect changes

* Adds tests for bugs changes in `vecdot`, `tensordot`

Author: ndgrigorian

dpctl/tensor/_linear_algebra_functions.py

@@ -90,8 +90,9 @@ def tensordot(x1, x2, axes=2):
             to `x2`. Both sequences must have equal length, and each axis
             `x1_axes[i]` for `x1` must have the same size as the respective
             axis `x2_axes[i]` for `x2`. Each sequence must consist of unique
-            non-negative integers that specify valid axes for each respective
-            array.
+            integers that specify valid axes for each respective array.
+            For example, if `x1` has rank `N`, a valid axis must reside on the
+            half-open interval `[-N, N)`.
     Returns:
         usm_ndarray:
             an array containing the tensor contraction whose shape consists of
@@ -154,11 +155,7 @@ def tensordot(x1, x2, axes=2):
         same_shapes = True
         for i in range(n_axes1):
             axis1 = axes1[i]
-            if axis1 < 0:
-                raise ValueError("`axes` must be non-negative")
             axis2 = axes2[i]
-            if axis2 < 0:
-                raise ValueError("`axes` must be non-negative")
             same_shapes = same_shapes and (x1_shape[axis1] == x2_shape[axis2])
         if not same_shapes:
             raise ValueError("shape mismatch in contracted `tensordot` axes")
@@ -314,12 +311,11 @@ def vecdot(x1, x2, axis=-1):
             axis. Input arrays should be of numeric type.
         axis (Optional[int]):
             axis over which to compute the dot product. The axis must
-            be an integer on the interval `[-N, N)`, where `N` is the
-            array rank of input arrays after broadcasting rules are
-            applied. If specified as a negative integer, the axis along
-            which dot product is performed is counted backward from
-            the last axes (that is `-1` refers to the last axis). By
-            default, dot product is computed over the last axis.
+            be an integer on the interval `[-N, -1]`, where `N` is
+            ``min(x1.ndim, x2.ndim)``. The axis along which dot product
+            is performed is counted backward from the last axes
+            (that is, `-1` refers to the last axis). By default,
+            dot product is computed over the last axis.
             Default: `-1`.
 
     Returns:
@@ -355,17 +351,19 @@ def vecdot(x1, x2, axis=-1):
     x2_nd = x2.ndim
     x1_shape = x1.shape
     x2_shape = x2.shape
+    if axis >= 0:
+        raise ValueError("`axis` must be negative")
+    axis = operator.index(axis)
+    x1_axis = normalize_axis_index(axis, x1_nd)
+    x2_axis = normalize_axis_index(axis, x2_nd)
+    if x1_shape[x1_axis] != x2_shape[x2_axis]:
+        raise ValueError(
+            "given axis must have the same shape for `x1` and `x2`"
+        )
     if x1_nd > x2_nd:
         x2_shape = (1,) * (x1_nd - x2_nd) + x2_shape
-        x2_nd = len(x2_shape)
     elif x2_nd > x1_nd:
         x1_shape = (1,) * (x2_nd - x1_nd) + x1_shape
-        x1_nd = len(x1_shape)
-    axis = normalize_axis_index(operator.index(axis), x1_nd)
-    if x1_shape[axis] != x2_shape[axis]:
-        raise ValueError(
-            "given axis must have the same shape for `x1` and `x2`"
-        )
     try:
         broadcast_sh = _broadcast_shape_impl(
             [
@@ -375,8 +373,10 @@ def vecdot(x1, x2, axis=-1):
         )
     except ValueError:
         raise ValueError("mismatch in `vecdot` dimensions")
+    broadcast_nd = len(broadcast_sh)
+    contracted_axis = normalize_axis_index(axis, broadcast_nd)
     res_sh = tuple(
-        [broadcast_sh[i] for i in range(len(broadcast_sh)) if i != axis]
+        [broadcast_sh[i] for i in range(broadcast_nd) if i != contracted_axis]
     )
     # type validation
     sycl_dev = exec_q.sycl_device
@@ -414,9 +414,8 @@ def vecdot(x1, x2, axis=-1):
             x1 = dpt.broadcast_to(x1, broadcast_sh)
         if x2.shape != broadcast_sh:
             x2 = dpt.broadcast_to(x2, broadcast_sh)
-        x1 = dpt.moveaxis(x1, axis, -1)
-        x2 = dpt.moveaxis(x2, axis, -1)
-
+        x1 = dpt.moveaxis(x1, contracted_axis, -1)
+        x2 = dpt.moveaxis(x2, contracted_axis, -1)
         out = dpt.empty(
             res_sh,
             dtype=res_dt,
@@ -427,7 +426,7 @@ def vecdot(x1, x2, axis=-1):
         ht_dot_ev, _ = tli._dot(
             x1=x1,
             x2=x2,
-            batch_dims=len(x1.shape[:-1]),
+            batch_dims=len(res_sh),
             x1_outer_dims=0,
             x2_outer_dims=0,
             inner_dims=1,
@@ -459,8 +458,8 @@ def vecdot(x1, x2, axis=-1):
             x1 = dpt.broadcast_to(x1, broadcast_sh)
         if buf2.shape != broadcast_sh:
             buf2 = dpt.broadcast_to(buf2, broadcast_sh)
-        x1 = dpt.moveaxis(x1, axis, -1)
-        buf2 = dpt.moveaxis(buf2, axis, -1)
+        x1 = dpt.moveaxis(x1, contracted_axis, -1)
+        buf2 = dpt.moveaxis(buf2, contracted_axis, -1)
         out = dpt.empty(
             res_sh,
             dtype=res_dt,
@@ -471,7 +470,7 @@ def vecdot(x1, x2, axis=-1):
         ht_dot_ev, _ = tli._dot(
             x1=x1,
             x2=buf2,
-            batch_dims=len(x1.shape[:-1]),
+            batch_dims=len(res_sh),
             x1_outer_dims=0,
             x2_outer_dims=0,
             inner_dims=1,
@@ -501,8 +500,8 @@ def vecdot(x1, x2, axis=-1):
             buf1 = dpt.broadcast_to(buf1, broadcast_sh)
         if x2.shape != broadcast_sh:
             x2 = dpt.broadcast_to(x2, broadcast_sh)
-        buf1 = dpt.moveaxis(buf1, axis, -1)
-        x2 = dpt.moveaxis(x2, axis, -1)
+        buf1 = dpt.moveaxis(buf1, contracted_axis, -1)
+        x2 = dpt.moveaxis(x2, contracted_axis, -1)
         out = dpt.empty(
             res_sh,
             dtype=res_dt,
@@ -513,7 +512,7 @@ def vecdot(x1, x2, axis=-1):
         ht_dot_ev, _ = tli._dot(
             x1=buf1,
             x2=x2,
-            batch_dims=len(x1.shape[:-1]),
+            batch_dims=len(res_sh),
             x1_outer_dims=0,
             x2_outer_dims=0,
             inner_dims=1,
@@ -548,8 +547,8 @@ def vecdot(x1, x2, axis=-1):
         buf1 = dpt.broadcast_to(buf1, broadcast_sh)
     if buf2.shape != broadcast_sh:
         buf2 = dpt.broadcast_to(buf2, broadcast_sh)
-    buf1 = dpt.moveaxis(buf1, axis, -1)
-    buf2 = dpt.moveaxis(buf2, axis, -1)
+    buf1 = dpt.moveaxis(buf1, contracted_axis, -1)
+    buf2 = dpt.moveaxis(buf2, contracted_axis, -1)
     out = dpt.empty(
         res_sh,
         dtype=res_dt,
@@ -560,7 +559,7 @@ def vecdot(x1, x2, axis=-1):
     ht_dot_ev, _ = tli._dot(
         x1=buf1,
         x2=buf2,
-        batch_dims=len(x1.shape[:-1]),
+        batch_dims=len(res_sh),
         x1_outer_dims=0,
         x2_outer_dims=0,
         inner_dims=1,

dpctl/tensor/libtensor/include/kernels/linalg_functions/dot_product.hpp

@@ -552,8 +552,8 @@ dot_product_contig_impl(sycl::queue &exec_q,
                 NoOpIndexerT, NoOpIndexerT>;
 
         const InputBatchIndexerT inp_batch_indexer{
-            0, static_cast<ssize_t>(reduction_nelems),
-            static_cast<ssize_t>(batches)};
+            0, static_cast<ssize_t>(batches),
+            static_cast<ssize_t>(reduction_nelems)};
         const InputOutputBatchIndexerT inp_out_batch_indexer{
             inp_batch_indexer, inp_batch_indexer, NoOpIndexerT{}};
         constexpr ReductionIndexerT reduction_indexer{NoOpIndexerT{},
@@ -588,8 +588,8 @@ dot_product_contig_impl(sycl::queue &exec_q,
                 NoOpIndexerT, NoOpIndexerT>;
 
         const InputBatchIndexerT inp_batch_indexer{
-            0, static_cast<ssize_t>(reduction_nelems),
-            static_cast<ssize_t>(batches)};
+            0, static_cast<ssize_t>(batches),
+            static_cast<ssize_t>(reduction_nelems)};
         const InputOutputBatchIndexerT inp_out_batch_indexer{
             inp_batch_indexer, inp_batch_indexer, NoOpIndexerT{}};
         constexpr ReductionIndexerT reduction_indexer{NoOpIndexerT{},
@@ -1174,8 +1174,8 @@ dot_product_contig_tree_impl(sycl::queue &exec_q,
                 NoOpIndexerT, NoOpIndexerT>;
 
         const InputBatchIndexerT inp_batch_indexer{
-            0, static_cast<ssize_t>(reduction_nelems),
-            static_cast<ssize_t>(batches)};
+            0, static_cast<ssize_t>(batches),
+            static_cast<ssize_t>(reduction_nelems)};
         const InputOutputBatchIndexerT inp_out_batch_indexer{
             inp_batch_indexer, inp_batch_indexer, NoOpIndexerT{}};
         constexpr ReductionIndexerT reduction_indexer{NoOpIndexerT{},
@@ -1212,8 +1212,8 @@ dot_product_contig_tree_impl(sycl::queue &exec_q,
                 NoOpIndexerT, NoOpIndexerT>;
 
         const InputBatchIndexerT inp_batch_indexer{
-            0, static_cast<ssize_t>(reduction_nelems),
-            static_cast<ssize_t>(batches)};
+            0, static_cast<ssize_t>(batches),
+            static_cast<ssize_t>(reduction_nelems)};
         const InputOutputBatchIndexerT inp_out_batch_indexer{
             inp_batch_indexer, inp_batch_indexer, NoOpIndexerT{}};
         constexpr ReductionIndexerT reduction_indexer{NoOpIndexerT{},
@@ -1280,8 +1280,8 @@ dot_product_contig_tree_impl(sycl::queue &exec_q,
                     NoOpIndexerT, NoOpIndexerT>;
 
             const InputBatchIndexerT inp_batch_indexer{
-                0, static_cast<ssize_t>(reduction_nelems),
-                static_cast<ssize_t>(batches)};
+                0, static_cast<ssize_t>(batches),
+                static_cast<ssize_t>(reduction_nelems)};
             const InputOutputBatchIndexerT inp_out_batch_indexer{
                 inp_batch_indexer, inp_batch_indexer, NoOpIndexerT{}};
             constexpr ReductionIndexerT reduction_indexer{NoOpIndexerT{},

dpctl/tensor/libtensor/include/kernels/linalg_functions/gemm.hpp

@@ -2100,7 +2100,7 @@ class GemmBatchNoAtomicFunctorThreadK
                 for (size_t vec_id = 1; vec_id < m_groups; ++vec_id) {
                     if (j + vec_id < m) {
                         res[total_offset + res_indexer(i * m + j + vec_id)] =
-                            local_sum[1];
+                            local_sum[vec_id];
                     }
                 }
             }

dpctl/tests/test_usm_ndarray_linalg.py

@@ -782,11 +782,16 @@ def test_tensordot_axes_errors():
     with pytest.raises(ValueError):
         dpt.tensordot(m1, m2, axes=-1)
 
-    with pytest.raises(ValueError):
-        dpt.tensordot(m1, m2, axes=((-1,), (1,)))
 
-    with pytest.raises(ValueError):
-        dpt.tensordot(m1, m2, axes=((1,), (-1,)))
+# tests for gh-1570
+def test_tensordot_gemm_small_k_m():
+    get_queue_or_skip()
+
+    x1 = dpt.asarray(1, dtype="i2")
+    x2 = dpt.asarray([0, 1, 0, 0], dtype="i2")
+
+    res = dpt.tensordot(x1, x2, axes=0)
+    assert dpt.all(x2 == res)
 
 
 @pytest.mark.parametrize("dtype", _numeric_types)
@@ -834,7 +839,7 @@ def test_vecdot_axis(dtype):
 
     v2 = dpt.ones((m1, n, m2), dtype=dtype)
 
-    r = dpt.vecdot(v1, v2, axis=1)
+    r = dpt.vecdot(v1, v2, axis=-2)
 
     assert r.shape == (
         m1,
@@ -864,7 +869,7 @@ def test_vecdot_strided(dtype):
         :, :n, ::-1
     ]
 
-    r = dpt.vecdot(v1, v2, axis=1)
+    r = dpt.vecdot(v1, v2, axis=-2)
 
     ref = sum(
         el1 * el2
@@ -903,6 +908,9 @@ def test_vector_arg_validation():
     with pytest.raises(ValueError):
         dpt.vecdot(v1, v2, axis=2)
 
+    with pytest.raises(ValueError):
+        dpt.vecdot(v1, v2, axis=-2)
+
     q = dpctl.SyclQueue(
         v2.sycl_context, v2.sycl_device, property="enable_profiling"
     )
@@ -946,3 +954,29 @@ def test_vecdot_type_promotion(dt1, dt2):
     assert r.shape == tuple()
     assert r.dtype == mul.dtype
     assert dpt.allclose(r, dpt.sum(mul, dtype=mul.dtype))
+
+
+def test_vecdot_broadcast_o1_buffer():
+    get_queue_or_skip()
+
+    v1 = dpt.arange(10, dtype="i2")
+    v2 = dpt.ones((5, 10), dtype="i4")
+
+    res1 = dpt.vecdot(v1, v2)
+    assert res1.shape == (5,)
+
+    res2 = dpt.vecdot(v2, v1)
+    assert res2.shape == (5,)
+
+
+def test_vecdot_contig_small():
+    get_queue_or_skip()
+
+    n = 1
+    for dt in [dpt.int16, dpt.int32, dpt.complex64]:
+        v1 = dpt.zeros((10, n), dtype=dt)
+        v2 = dpt.ones_like(v1, dtype=dt)
+        v1[-1] = 1
+        res = dpt.vecdot(v1, v2)
+        assert dpt.all(res[:-1] == 0)
+        assert res[-1] == n