CustomArray

CustomArray#

class brainunit.CustomArray[source]#

A custom array wrapper providing comprehensive array operations and cross-framework compatibility.

CustomArray is a mix-in class that delegates every operation to its data attribute. Subclasses must provide a data property (or attribute) that returns the underlying array. The class exposes NumPy-style methods, PyTorch-style convenience methods, and full interoperability with JAX transformations (jit, grad, vmap).

data#

The underlying array storage. Subclasses are responsible for providing this attribute (e.g. via a property backed by some internal state).

Type:: array-like

NumPy-style methods: all, any, argmax, argmin, argsort, astype, clip, copy, cumsum, cumprod, diagonal, dot, flatten, max, mean, min, nonzero, prod, ravel, repeat, reshape, round, squeeze, std, sum, swapaxes, take, tolist, trace, transpose, var

PyTorch-style methods: unsqueeze, expand, expand_as, clamp, clone, zero_, bool, int, long, half, float, double

Conversion methods: to_numpy, to_jax, numpy

Trigonometric methods: sin, cos, tan, sinh, cosh, tanh, arcsin, arccos, arctan (and in-place _ variants)

Examples

CustomArray is designed to be used as a mix-in. A minimal standalone subclass needs only a data attribute and JAX pytree registration:

import jax
import numpy as np
from saiunit import CustomArray

@jax.tree_util.register_pytree_node_class
class SimpleArray(CustomArray):
    def __init__(self, value):
        self.data = value

    def tree_flatten(self):
        return (self.data,), None

    @classmethod
    def tree_unflatten(cls, aux_data, flat_contents):
        return cls(*flat_contents)

Basic properties and arithmetic:

arr = SimpleArray(np.array([1.0, 2.0, 3.0]))
arr.shape   # (3,)
arr.ndim    # 1
arr + 10    # array([11., 12., 13.])
arr ** 2    # array([1., 4., 9.])

Statistical operations:

arr = SimpleArray(np.array([1.0, 2.0, 3.0, 4.0, 5.0]))
arr.mean()  # 3.0
arr.std()   # ~1.414
arr.sum()   # 15.0

Array manipulation:

matrix = SimpleArray(np.array([[1, 2, 3], [4, 5, 6]]))
matrix.T           # transposed (3, 2) array
matrix.reshape(6)  # array([1, 2, 3, 4, 5, 6])
matrix.flatten()   # array([1, 2, 3, 4, 5, 6])

JAX compatibility (jit, grad):

import jax
import jax.numpy as jnp

arr = SimpleArray(jnp.array([1.0, 2.0, 3.0]))

@jax.jit
def square(x):
    return x * x

square(arr)  # Array([1., 4., 9.], ...)

Notes

This class uses duck typing and delegates operations to the underlying data attribute.
In-place operations (+=, -=, etc.) modify data directly and return self.
Most methods return the raw underlying array type, not a new CustomArray instance.
Thread safety depends on the underlying array implementation.
JAX transformations (jit, grad, vmap) work seamlessly when the subclass is registered as a JAX pytree.

CustomArray

Contents

CustomArray#