[P] dataclass_array: Dataclasses which can be manipulated like numpy arrays (batched, reshaped, sliced,...) (for TF, Jax, Numpy,...)

Available at: [https://github.com/google-research/dataclass\_array](https://github.com/google-research/dataclass_array)

\`dataclass\_array\` allow to have structured data that can be of arbitrary batch shape. For example:

Defining a dataclass array:

@dataclasses.dataclass(frozen=True)
class Ray(dca.DataclassArray):
pos: FloatArray['*batch_shape 3']
dir: FloatArray['*batch_shape 3']

Dataclass array can then be manipulated as if they were ndarray, while keeping the internal semantic structure

rays = camera.rays() # Returns `Ray` with shape `(h, w)`

rays.shape == (h, w)
rays.pos.shape == (h, w, 3) # Individual ndarray fields accessible

rays = rays.reshape('h w -> w h') # …

arrays jax machinelearning numpy