Is there a way to have record array in JAX like in Numpy?

[a11to1n3]

Hello community,

in my project I am converting a numpy-driven library into jax, in which I find it is really convenient to have record array which can store object members (numpy.lib.recfunctions). I would like to ask if there is a way to do such similar array in jax?

[jakevdp]

No, JAX does not support compound dtypes, so there is no way to natively implement a record array in JAX. If you want this sort of functionality, I'd recommend a dataclass with multiple array buffers to store the multiple elements.