eriknw · Sep 24, 2026 · Jul 31, 2026 · Aug 6, 2026 · Sep 23, 2026 · Sep 23, 2026
diff --git a/docs/user_guide/udt.rst b/docs/user_guide/udt.rst
@@ -32,7 +32,9 @@ A UDT is any ``numpy.dtype`` you register with python-graphblas. There are three
     point3 = dtypes.register_anonymous(np.dtype((np.float64, (3,))), "Point3")
 
 Multi-dimensional shapes work too (``np.dtype((np.float64, (2, 4)))``); the
-layout is flattened row-major in C.
+layout is flattened row-major in C. An array of arrays has the same layout, so
+it registers as the same UDT: ``np.dtype((point3.np_type, (2,)))`` is
+``FP64[2, 3]``. Record fields that nest arrays are flattened the same way.
 
 **Dataclass UDTs** are record UDTs derived from a ``@dataclass``:
 
@@ -68,7 +70,11 @@ Field type rules:
 
 - Numeric scalar types (``int``, ``float``, ``bool``, ``complex``, the
   corresponding numpy scalar types) are supported.
-- Strings, Python objects, and nested UDTs are not.
+- Fixed-shape arrays (``("pos", np.float64, (3,))``) and records (a nested
+  struct) are supported as fields, and built-in operators lift through them.
+  In the dict and dataclass forms, write an array field as ``"FP64[3]"``; the
+  dict form also takes an existing UDT, such as ``{"id": int, "pos": point3}``.
+- Strings and Python objects are not supported.
 - Dataclass annotation strings (e.g., ``"int"``) resolve through
   ``lookup_dtype``.
 
@@ -173,6 +179,26 @@ For nested record UDTs the tuple is *flat over the leaves*. Given
 return ``(id, x, y)``, not ``(id, (x, y))``. Returning an existing record value
 (e.g., one of the inputs) is also fine and preserves the nested shape.
 
+For array UDTs each operand arrives as a numpy view of that element's values, in
+the UDT's declared shape: a ``np.dtype((np.float64, (2, 4)))`` UDT hands the UDF
+a 2-by-4 array, indexable as ``x[i, j]``. Array expressions work as written, and
+the UDF may return one of its operands or build a new array of the same shape:
+
+.. code-block:: python
+
+    def midpoint(x, y):
+        return (x + y) / 2
+
+    op = binary.register_new("midpoint", midpoint, is_udt=True)
+
+    a = Vector(point3, size=1)
+    a[0] = [0.0, 2.0, 4.0]
+    b = Vector(point3, size=1)
+    b[0] = [10.0, 20.0, 30.0]
+
+    c = a.ewise_mult(b, op[point3]).new()
+    # c[0] = [5.0, 11.0, 17.0]
+
 If your UDF references a field that doesn't exist, or returns the wrong arity,
 you'll get a ``UdfParseError`` with the actionable diagnostic line surfaced
 from Numba's typing pass instead of a 200-line traceback.
@@ -216,6 +242,8 @@ JIT is skipped when:
   falls through to the Numba cfunc path.
 - A field type isn't in the numpy-to-C map (rare; the standard numeric
   scalar types all map).
+- A record has an array-valued field. Built-in operators still lift to it
+  through the Numba cfunc.
 - The numpy layout doesn't match what a C compiler would produce. The most
   common case is a packed record with mixed-width fields (e.g.,
   ``np.dtype([("a", int32), ("b", float64)])`` without ``align=True``).

diff --git a/graphblas/core/dtypes.py b/graphblas/core/dtypes.py
@@ -10,6 +10,7 @@
 
 if _has_numba:
     import numba
+    from numba.np.numpy_support import as_dtype as _numba_as_dtype
 
 # Default assumption unless FC32/FC64 are found in lib
 _supports_complex = hasattr(lib, "GrB_FC64") or hasattr(lib, "GxB_FC64")
@@ -200,6 +201,10 @@ def register_anonymous(dtype, name=None):
     name; the SuiteSparse-side C name set at first registration does not
     change. See ``DataType.jit_c_name``.
 
+    Nested subarrays are flattened first, so ``np.dtype((np.dtype((np.float64,
+    (3,))), (2,)))``, two ``FP64[3]`` elements, registers as ``FP64[2, 3]`` and
+    returns the same ``DataType``. This applies inside records too.
+
     Parameters
     ----------
     dtype : np.dtype | dict | str | @dataclass
@@ -259,6 +264,7 @@ def register_anonymous(dtype, name=None):
             dtype = np.dtype((base_dtype.np_type, shape))
         else:
             raise
+    dtype = _flatten_subarrays(dtype)
     if dtype in _registry:
         # Always use the same object, but use the latest name. The
         # Python-side ``rv.name`` updates here, but the SuiteSparse-side
@@ -318,9 +324,6 @@ def register_anonymous(dtype, name=None):
     rv = DataType(name, gb_obj, None, f"uint8_t[{dtype.itemsize}]", numba_type, dtype)
     _registry[gb_obj] = rv
     _registry[dtype] = rv
-    if _has_numba:
-        _registry[numba_type] = rv
-        _registry[numba_type.name] = rv
     # Set JIT C type definition so SuiteSparse can JIT-compile kernels for this UDT.
     _set_udt_jit_c_definition(rv)
     return rv
@@ -532,6 +535,12 @@ def lookup_dtype(key, value=None):
         return _registry[key]
     except (KeyError, TypeError):
         pass
+    if _has_numba and isinstance(key, (numba.types.Record, numba.types.NestedArray)):
+        # A UDT's Numba type is not a registry key. Convert it to its numpy
+        # dtype, which is already flat, and look that up, so a layout no UDT
+        # has registered yet (e.g., a record's array field returned by a UDF)
+        # resolves the same way regardless of what else is registered.
+        return lookup_dtype(_numba_as_dtype(key))
     if value is not None and hasattr(value, "dtype") and value.dtype in _registry:
         return _registry[value.dtype]
     # np.dtype(x) accepts some weird values; we may want to guard against some
@@ -585,6 +594,61 @@ def _default_name(dtype):
     return repr(dtype)
 
 
+def _flatten_subarrays(np_type):
+    """Collapse each run of nested subarray dtypes in ``np_type`` into one subarray.
+
+    numpy keeps ``np.dtype((np.dtype((np.float64, (3,))), (2,)))`` layered, as
+    two ``FP64[3]`` elements, and compares it unequal to the flat
+    ``np.dtype((np.float64, (2, 3)))``. Nothing else tells them apart: the bytes
+    are the same, an array of either expands to a ``(..., 2, 3)`` float64 array,
+    and Numba types both as ``nestedarray(float64, (2, 3))``. Registering the flat
+    form gives each layout one DataType. Records are rebuilt only when a field
+    changes, keeping their offsets, titles, itemsize, and alignment flag.
+    """
+    if np_type.subdtype is not None:
+        base, shape = np_type.subdtype
+        flat_base = _flatten_subarrays(base)
+        if flat_base.subdtype is not None:
+            flat_base, inner_shape = flat_base.subdtype
+            shape += inner_shape
+        elif flat_base is base:
+            return np_type
+        return np.dtype((flat_base, shape))
+    if np_type.names is not None:
+        fields = [np_type.fields[name] for name in np_type.names]
+        formats = [_flatten_subarrays(field[0]) for field in fields]
+        if all(new is field[0] for new, field in zip(formats, fields, strict=True)):
+            return np_type
+        spec = {
+            "names": list(np_type.names),
+            "formats": formats,
+            "offsets": [field[1] for field in fields],
+            "itemsize": np_type.itemsize,
+        }
+        if any(len(field) > 2 for field in fields):
+            spec["titles"] = [field[2] if len(field) > 2 else None for field in fields]
+        return np.dtype(spec, align=np_type.isalignedstruct)
+    return np_type
+
+
+def _view_if_same_layout(values, np_type):
+    """Return ``values`` viewed as ``np_type`` if the two differ only in subarray nesting.
+
+    A record whose field nests subarrays has a different dtype from its
+    registered, flattened form, and numpy casts between the two wrongly in
+    either direction: ``astype``, ``np.array(..., dtype=)``, and item assignment
+    all scramble the field's values. The bytes already match, so a view is exact.
+    """
+    if (
+        isinstance(values, (np.ndarray, np.void))
+        and values.dtype != np_type
+        and values.dtype.itemsize == np_type.itemsize
+        and _flatten_subarrays(values.dtype) == np_type
+    ):
+        return values.view(np_type)
+    return values
+
+
 def _dtype_to_string(dtype):
     """Convert a dtype to a string that can be safely evaluated to recreate the dtype.
 

diff --git a/graphblas/core/operator/base.py b/graphblas/core/operator/base.py
@@ -272,7 +272,9 @@ def _resolve_udt_return_type(numba_ret_type, *dtypes):
 
         When a UDF returns a tuple, Numba infers ``Tuple(...)`` rather than a
         Record type. Match by field count, preferring a candidate whose field
-        types align with the Tuple's element types.
+        types align with the Tuple's element types. An array the UDF builds
+        (``x + y``) is typed without its extents, so it takes the type of the
+        array-UDT operand with the same element type and rank.
         """
         try:
             return lookup_dtype(numba_ret_type)
@@ -329,17 +331,80 @@ def _leaves(d):
                     f"shape {shape}. Return a numpy array (e.g., ``np.array(...)``) or a "
                     f"scalar; tuple returns are only matched to record UDTs."
                 )
+        elif isinstance(numba_ret_type, numba.core.types.Array):
+            # Numba types an array the UDF builds as a plain ``Array``: element
+            # type and rank, but no extents, which only an operand can supply.
+            # (An operand returned as-is keeps its ``NestedArray`` type, extents
+            # and all, and ``lookup_dtype`` resolved it above.)
+            array_inputs = list(
+                dict.fromkeys(  # dedup: ``x`` and ``y`` are often the same UDT
+                    d for d in dtypes if d._is_udt and d.np_type.subdtype is not None
+                )
+            )
+            if isinstance(numba_ret_type, numba.core.types.NestedArray):
+                # Extents survived, so ``lookup_dtype`` above named a UDT for
+                # them unless SuiteSparse refused to create it: builds without
+                # variable-length arrays cap a UDT at 128 bytes before SS 9.0
+                # and 1024 from then on. Matching those extents to a same-rank
+                # operand would hand SuiteSparse an element of the wrong size.
+                shaped_like = []
+            else:
+                shaped_like = [
+                    d
+                    for d in array_inputs
+                    if d.numba_type.dtype == numba_ret_type.dtype
+                    and d.numba_type.ndim == numba_ret_type.ndim
+                ]
+            if len(shaped_like) == 1:
+                return shaped_like[0]
+            if shaped_like:
+                # Operands that differ only in length are indistinguishable
+                # here. Guessing would hand SuiteSparse an element of the
+                # wrong size, so say so.
+                raise UdfParseError(
+                    f"UDT UDF returned {numba_ret_type!r}, which matches more than one "
+                    f"input array UDT ({', '.join(str(d) for d in shaped_like)}). "
+                    f"Return one of the operands, or make the operands the same type."
+                )
+            # An array UDT went in and an array came out, but not one that
+            # fits: name the mismatch rather than fall through to the generic
+            # "unsupported type", whose advice the user already followed.
+            if array_inputs:
+                described = "; ".join(
+                    f"{d} elements are {d.numba_type.dtype} with shape {d.numba_type.shape}"
+                    for d in array_inputs
+                )
+                raise UdfParseError(
+                    f"UDT UDF returned {numba_ret_type!r}, which matches no input array "
+                    f"UDT: {described}. Return an array matching one of those, or one of "
+                    f"the operands."
+                )
         raise UdfParseError(
             f"UDT UDF returned an unsupported type {numba_ret_type!r}. "
             f"Return a scalar, a tuple matching a record UDT's fields, or a numpy array "
             f"matching an array UDT's shape."
         )
 
+    def _array_udt_view(dtype):
+        """Return ``(base_element_numba_type, shape)`` for an array UDT.
+
+        Array UDTs are addressed as a ``carray`` over their base elements, in
+        the UDT's declared shape, rather than as Numba's ``NestedArray``. Numba
+        models a ``NestedArray`` *value* as an array descriptor (data pointer,
+        shape, strides, ...), so loading or storing one through a ``CPointer``
+        moves the descriptor rather than the element payload, corrupting
+        whatever follows it (and overrunning the element outright once the
+        descriptor is the wider of the two).
+        """
+        nested = dtype.numba_type
+        return nested.dtype, nested.shape
+
     def _input_operand(dtype, var):
         """Return ``(setup_line, deref_expr, ptr_arg_type)`` for one input operand.
 
-        - ``setup_line`` is the optional ``var = numba.carray(var_ptr, 1)`` line
-          to add to the wrapper body (empty for non-record cases).
+        - ``setup_line`` is the optional ``var = numba.carray(var_ptr, n)`` line
+          to add to the wrapper body, where ``n`` is 1 for a record UDT and the
+          element's shape for an array UDT (empty for non-UDT cases).
         - ``deref_expr`` is the value to pass to ``numba_func`` for this operand.
         - ``ptr_arg_type`` is the Numba ``CPointer(...)`` type for the wrapper signature.
         """
@@ -351,8 +416,12 @@ def _input_operand(dtype, var):
                     f"{var}[0]",
                     nt.CPointer(dtype.numba_type),
                 )
-            # Array UDT: pass the raw pointer. The UDF carrays it if it wants.
-            return "", f"{var}_ptr", nt.CPointer(dtype.numba_type)
+            base_numba, shape = _array_udt_view(dtype)
+            return (
+                f"    {var} = numba.carray({var}_ptr, {shape})\n",
+                var,
+                nt.CPointer(base_numba),
+            )
         if dtype == BOOL:
             # Numba can't compile bool ptrs (numba/numba#5395); expose them
             # as int8 and cast on deref.
@@ -364,9 +433,11 @@ def _output_handler(return_type, numba_ret_type):
         """Return ``(setup_line, ret_ptr_type, write_kind, write_info)``.
 
         ``write_kind`` is ``"record_fields"`` when the UDF returns a Tuple to be
-        unpacked into a record output, otherwise ``"direct"``.
-        ``write_info`` is the tuple-unpack field tuple for ``"record_fields"``,
-        or a ``(BL, BR, zname)`` 3-tuple for ``"direct"``.
+        unpacked into a record output, ``"array_elements"`` for an array UDT
+        output, otherwise ``"direct"``. ``write_info`` is one
+        ``(access_path, is_array)`` pair per leaf for ``"record_fields"`` and a
+        ``(BL, BR, zname)`` 3-tuple for ``"direct"``; ``"array_elements"`` needs
+        none, since ``setup_line`` already binds ``z`` to the right shape.
         """
         nt = numba.types
         ztype = INT8 if return_type == BOOL else return_type
@@ -378,18 +449,21 @@ def _output_handler(return_type, numba_ret_type):
             and return_type._is_udt
             and return_type.np_type.names is not None
         ):
-            # ``write_info`` is the list of Python access paths to each leaf
-            # field (``"['a']"`` for flat, ``"['outer']['inner_a']"`` for
-            # nested). The wrapper iterates these to write the flat tuple
-            # ``_result`` back leaf-by-leaf.
+            # ``write_info`` pairs each leaf field's Python access path
+            # (``"['a']"`` for flat, ``"['outer']['inner_a']"`` for nested)
+            # with whether that leaf is array-typed. The wrapper iterates
+            # these to write the flat tuple ``_result`` back leaf-by-leaf.
             from .udt_utils import _iter_record_leaves
 
-            leaf_python_paths = tuple(py for py, _c, _d in _iter_record_leaves(return_type.np_type))
+            leaves = tuple(
+                (py, d.subdtype is not None)
+                for py, _c, d in _iter_record_leaves(return_type.np_type)
+            )
             return (
                 "    z = numba.carray(z_ptr, 1)\n",
                 ret_ptr_type,
                 "record_fields",
-                leaf_python_paths,
+                leaves,
             )
         if return_type._is_udt:
             if return_type.np_type.subdtype is None:
@@ -399,9 +473,13 @@ def _output_handler(return_type, numba_ret_type):
                     "direct",
                     ("", "", "z[0]"),
                 )
-            # Array UDT: write via z_ptr[0]. The UDF receives the carray and
-            # writes in place.
-            return "", ret_ptr_type, "direct", ("", "[0]", "z_ptr[0]")
+            base_numba, shape = _array_udt_view(return_type)
+            return (
+                f"    z = numba.carray(z_ptr, {shape})\n",
+                nt.CPointer(base_numba),
+                "array_elements",
+                None,
+            )
         if return_type == BOOL:
             return "", ret_ptr_type, "direct", ("bool(", ")", "z_ptr[0]")
         return "", ret_ptr_type, "direct", ("", "", "z_ptr[0]")
@@ -410,11 +488,29 @@ def _compose_wrapper_body(zkind, zinfo, signature_line, body_setup, call_expr):
         """Assemble the Python source for a UDT cfunc wrapper.
 
         For record returns, the wrapper writes leaf-by-leaf; Numba does not
-        compile a nested-tuple assignment to an outer record field.
+        compile a nested-tuple assignment to an outer record field. Array
+        returns and array-typed record leaves slice-assign into the caller's
+        buffer, so a wrong-shape return cannot overrun it. The error does not reach
+        the caller either: the resulting ``ValueError`` is raised inside a
+        cfunc, which Numba prints and swallows, so the write stops there and
+        every leaf from that point on keeps whatever SuiteSparse had in the
+        buffer.
         """
+        if zkind == "array_elements":
+            return (
+                f"{signature_line}\n"
+                f"{body_setup}"
+                f"    _result = {call_expr}\n"
+                f"    z[:] = _result\n"
+            )
         if zkind == "record_fields":
+            # Array-typed leaves slice-assign. Numba's record-field setitem
+            # copies the destination's extent regardless of the source's, so a
+            # short source is read past its end and a long one is silently
+            # truncated; ``[:]`` makes both a shape error instead.
             field_assigns = "".join(
-                f"    z[0]{path} = _result[{i}]\n" for i, path in enumerate(zinfo)
+                f"    z[0]{path}{'[:]' if is_array else ''} = _result[{i}]\n"
+                for i, (path, is_array) in enumerate(zinfo)
             )
             return (
                 f"{signature_line}\n"

diff --git a/graphblas/core/operator/udt_utils.py b/graphblas/core/operator/udt_utils.py
@@ -349,9 +349,11 @@ def _iter_record_leaves(np_type, python_prefix="", c_prefix=""):
     ``("['outer']['inner_a']", "outer.inner_a", float64_dtype)``.
 
     Only record-in-record nesting is recognised. Array-typed fields are
-    yielded as a single leaf (the codegen treats them opaquely; arithmetic
-    on an array-typed sub-field isn't supported by either the cfunc or the
-    JIT path today).
+    yielded as a single leaf, which the Numba path handles: the generated
+    expression operates on the whole sub-array and the wrapper slice-assigns
+    it back. The JIT path never sees such a record, because a subarray dtype
+    has no entry in ``NP_TO_C_TYPES`` and ``_udt_c_typedef`` bails out, so
+    the op keeps the cfunc for every call.
     """
     for name in np_type.names:
         field_dtype = np_type.fields[name][0]