[Type] Tensor 3: Add qd.tensor_vec / qd.tensor_mat

docs/source/user_guide/flexible_tensors.md

@@ -77,6 +77,31 @@ qd.tensor(qd.f32, shape=(3,), backend="field")  # ValueError
 Integer values (`0`, `1`) are accepted because `Backend` is an `IntEnum`, but
 prefer the named members for clarity at call sites.
 
-Subsequent releases will add `qd.tensor_vec` / `qd.tensor_mat` for compound
-element types, a `qd.tensor_annotation(backend)` helper for kernel argument
-typing, and a `layout=` keyword for per-tensor physical-memory layout.
+## Vector and matrix tensors
+
+For tensors whose elements are vectors or matrices, use `qd.tensor_vec` and
+`qd.tensor_mat`. They dispatch over `qd.Vector.field` / `qd.Vector.ndarray`
+and `qd.Matrix.field` / `qd.Matrix.ndarray` respectively, with the same
+`backend=` keyword:
+
+```python
+import quadrants as qd
+
+qd.init(arch=qd.x64)
+
+# A 1-D tensor of 4 length-3 vectors, on the field backend (default).
+v = qd.tensor_vec(3, qd.f32, shape=(4,))
+
+# Same shape, on the ndarray backend.
+u = qd.tensor_vec(3, qd.f32, shape=(4,), backend=qd.Backend.NDARRAY)
+
+# A 1-D tensor of 3 (2x2) matrices, on the field backend.
+m = qd.tensor_mat(2, 2, qd.f32, shape=(3,))
+```
+
+These match the existing `qd.Vector.*` / `qd.Matrix.*` factories one-for-one;
+`qd.tensor_vec` / `qd.tensor_mat` simply add the per-tensor `backend=` knob.
+
+Subsequent releases will add a `qd.tensor_annotation(backend)` helper for
+kernel argument typing, and a `layout=` keyword for per-tensor physical-memory
+layout.

python/quadrants/__init__.py

@@ -60,6 +60,8 @@ def __getattr__(attr):
     "math",
     "sparse",
     "tensor",
+    "tensor_mat",
+    "tensor_vec",
     "tools",
     "types",
 ]

python/quadrants/_flexible.py

@@ -8,7 +8,7 @@
 
 from enum import IntEnum
 
-__all__ = ["Backend", "tensor"]
+__all__ = ["Backend", "tensor", "tensor_mat", "tensor_vec"]
 
 
 class Backend(IntEnum):
@@ -81,3 +81,72 @@ def tensor(dtype, shape, *, backend=Backend.FIELD, **kwargs):
     if backend is Backend.NDARRAY:
         return impl.ndarray(dtype, shape, **kwargs)
     raise AssertionError(f"unhandled Backend member: {backend!r}")
+
+
+def tensor_vec(n, dtype, shape, *, backend=Backend.FIELD, **kwargs):
+    """Allocate a tensor whose elements are length-``n`` vectors.
+
+    Dispatcher over ``qd.Vector.field`` and ``qd.Vector.ndarray`` selected by
+    the ``backend=`` keyword.
+
+    Args:
+        n (int): Length of each vector element.
+        dtype: Element data type (e.g. ``qd.f32``).
+        shape: Shape of the tensor (excluding the vector dimension) as an
+            ``int`` or tuple of ``int``.
+        backend (Backend, optional): Storage backend. Defaults to
+            :attr:`Backend.FIELD`.
+        **kwargs: Forwarded verbatim to the underlying ``qd.Vector.field`` /
+            ``qd.Vector.ndarray`` call. ``qd.Vector.ndarray`` does not accept
+            extra keyword arguments today.
+
+    Example::
+
+        >>> import quadrants as qd
+        >>> qd.init(arch=qd.x64)
+        >>> v = qd.tensor_vec(3, qd.f32, shape=(4,))
+        >>> u = qd.tensor_vec(3, qd.f32, shape=(4,), backend=qd.Backend.NDARRAY)
+    """
+    backend = _coerce_backend(backend)
+    from quadrants.lang.matrix import Vector  # late import
+
+    if backend is Backend.FIELD:
+        return Vector.field(n, dtype, shape, **kwargs)
+    if backend is Backend.NDARRAY:
+        return Vector.ndarray(n, dtype, shape, **kwargs)
+    raise AssertionError(f"unhandled Backend member: {backend!r}")
+
+
+def tensor_mat(n, m, dtype, shape, *, backend=Backend.FIELD, **kwargs):
+    """Allocate a tensor whose elements are ``n``-by-``m`` matrices.
+
+    Dispatcher over ``qd.Matrix.field`` and ``qd.Matrix.ndarray`` selected by
+    the ``backend=`` keyword.
+
+    Args:
+        n (int): Number of rows of each matrix element.
+        m (int): Number of columns of each matrix element.
+        dtype: Element data type (e.g. ``qd.f32``).
+        shape: Shape of the tensor (excluding the matrix dimensions) as an
+            ``int`` or tuple of ``int``.
+        backend (Backend, optional): Storage backend. Defaults to
+            :attr:`Backend.FIELD`.
+        **kwargs: Forwarded verbatim to the underlying ``qd.Matrix.field`` /
+            ``qd.Matrix.ndarray`` call. ``qd.Matrix.ndarray`` does not accept
+            extra keyword arguments today.
+
+    Example::
+
+        >>> import quadrants as qd
+        >>> qd.init(arch=qd.x64)
+        >>> a = qd.tensor_mat(2, 3, qd.f32, shape=(4,))
+        >>> b = qd.tensor_mat(2, 3, qd.f32, shape=(4,), backend=qd.Backend.NDARRAY)
+    """
+    backend = _coerce_backend(backend)
+    from quadrants.lang.matrix import Matrix  # late import
+
+    if backend is Backend.FIELD:
+        return Matrix.field(n, m, dtype, shape, **kwargs)
+    if backend is Backend.NDARRAY:
+        return Matrix.ndarray(n, m, dtype, shape, **kwargs)
+    raise AssertionError(f"unhandled Backend member: {backend!r}")

tests/python/test_flexible_factory_vec_mat.py

@@ -0,0 +1,136 @@
+"""Tests for ``qd.tensor_vec`` / ``qd.tensor_mat`` (PR 3)."""
+
+import pytest
+
+import quadrants as qd
+
+from tests import test_utils
+
+
+# ----------------------------------------------------------------------------
+# qd.tensor_vec
+# ----------------------------------------------------------------------------
+
+
+@test_utils.test(arch=qd.cpu)
+def test_tensor_vec_default_backend_matches_vector_field():
+    a = qd.tensor_vec(3, qd.f32, shape=(4,))
+    b = qd.Vector.field(3, qd.f32, shape=(4,))
+    assert type(a) is type(b)
+    assert a.shape == b.shape == (4,)
+
+
+@test_utils.test(arch=qd.cpu)
+def test_tensor_vec_field_explicit():
+    a = qd.tensor_vec(3, qd.f32, shape=(4,), backend=qd.Backend.FIELD)
+    ref = qd.Vector.field(3, qd.f32, shape=(4,))
+    assert type(a) is type(ref)
+
+
+@test_utils.test(arch=qd.cpu)
+def test_tensor_vec_ndarray_matches_vector_ndarray():
+    a = qd.tensor_vec(3, qd.f32, shape=(4,), backend=qd.Backend.NDARRAY)
+    ref = qd.Vector.ndarray(3, qd.f32, shape=(4,))
+    assert type(a) is type(ref)
+    assert a.shape == ref.shape == (4,)
+
+
+@test_utils.test(arch=qd.cpu)
+def test_tensor_vec_invalid_backend_raises():
+    with pytest.raises(ValueError, match="backend="):
+        qd.tensor_vec(3, qd.f32, shape=(4,), backend="oops")
+
+
+@test_utils.test(arch=qd.cpu)
+def test_tensor_vec_kernel_roundtrip_field():
+    v = qd.tensor_vec(3, qd.f32, shape=(4,))
+
+    @qd.kernel
+    def fill(x: qd.template()):
+        for i in range(4):
+            for j in qd.static(range(3)):
+                x[i][j] = i * 10.0 + j
+
+    fill(v)
+    arr = v.to_numpy()
+    assert arr.shape == (4, 3)
+    assert arr[2, 1] == 21.0
+
+
+@test_utils.test(arch=qd.cpu)
+def test_tensor_vec_kernel_roundtrip_ndarray():
+    v = qd.tensor_vec(3, qd.f32, shape=(4,), backend=qd.Backend.NDARRAY)
+
+    @qd.kernel
+    def fill(x: qd.types.ndarray()):
+        for i in range(4):
+            for j in qd.static(range(3)):
+                x[i][j] = i * 10.0 + j
+
+    fill(v)
+    arr = v.to_numpy()
+    assert arr.shape == (4, 3)
+    assert arr[2, 1] == 21.0
+
+
+# ----------------------------------------------------------------------------
+# qd.tensor_mat
+# ----------------------------------------------------------------------------
+
+
+@test_utils.test(arch=qd.cpu)
+def test_tensor_mat_default_backend_matches_matrix_field():
+    a = qd.tensor_mat(2, 3, qd.f32, shape=(4,))
+    b = qd.Matrix.field(2, 3, qd.f32, shape=(4,))
+    assert type(a) is type(b)
+    assert a.shape == b.shape == (4,)
+
+
+@test_utils.test(arch=qd.cpu)
+def test_tensor_mat_ndarray_matches_matrix_ndarray():
+    a = qd.tensor_mat(2, 3, qd.f32, shape=(4,), backend=qd.Backend.NDARRAY)
+    ref = qd.Matrix.ndarray(2, 3, qd.f32, shape=(4,))
+    assert type(a) is type(ref)
+    assert a.shape == ref.shape == (4,)
+
+
+@test_utils.test(arch=qd.cpu)
+def test_tensor_mat_invalid_backend_raises():
+    with pytest.raises(ValueError, match="backend="):
+        qd.tensor_mat(2, 3, qd.f32, shape=(4,), backend=99)
+
+
+@test_utils.test(arch=qd.cpu)
+def test_tensor_mat_kernel_roundtrip_field():
+    m = qd.tensor_mat(2, 2, qd.f32, shape=(3,))
+
+    @qd.kernel
+    def fill(x: qd.template()):
+        for i in range(3):
+            for r in qd.static(range(2)):
+                for c in qd.static(range(2)):
+                    x[i][r, c] = i * 100.0 + r * 10.0 + c
+
+    fill(m)
+    arr = m.to_numpy()
+    assert arr.shape == (3, 2, 2)
+    assert arr[1, 0, 1] == 101.0
+    assert arr[2, 1, 0] == 210.0
+
+
+@test_utils.test(arch=qd.cpu)
+def test_tensor_mat_kernel_roundtrip_ndarray():
+    m = qd.tensor_mat(2, 2, qd.f32, shape=(3,), backend=qd.Backend.NDARRAY)
+
+    @qd.kernel
+    def fill(x: qd.types.ndarray()):
+        for i in range(3):
+            for r in qd.static(range(2)):
+                for c in qd.static(range(2)):
+                    x[i][r, c] = i * 100.0 + r * 10.0 + c
+
+    fill(m)
+    arr = m.to_numpy()
+    assert arr.shape == (3, 2, 2)
+    assert arr[1, 0, 1] == 101.0
+    assert arr[2, 1, 0] == 210.0