from docarray import BaseDoc
from docarray.typing import JaxArray
import jax.numpy as jnp


class MyDoc(BaseDoc):
    arr: JaxArray
    image_arr: JaxArray[3, 224, 224]
    square_crop: JaxArray[3, 'x', 'x']
    random_image: JaxArray[3, ...]  # first dimension is fixed, can have arbitrary shape


# create a document with tensors
doc = MyDoc(
    arr=jnp.zeros((128,)),
    image_arr=jnp.zeros((3, 224, 224)),
    square_crop=jnp.zeros((3, 64, 64)),
    random_image=jnp.zeros((3, 128, 256)),
)
assert doc.image_arr.shape == (3, 224, 224)

# automatic shape conversion
doc = MyDoc(
    arr=np.zeros((128,)),
    image_arr=np.zeros((224, 224, 3)),  # will reshape to (3, 224, 224)
    square_crop=np.zeros((3, 128, 128)),
    random_image=np.zeros((3, 64, 128)),
)
assert doc.image_arr.shape == (3, 224, 224)