Failed to Get a Single Element of Tensor

[daskol]

The following snippet fails with ttax of version 0.0.3 installed from PyPI.

>>> import jax, ttax
>>> key = jax.random.PRNGKey(42)
>>> a = ttax.random.tensor(key, (2, ) * 3, 2)
>>> a[0, 0, 0]
File ~/.local/lib/python3.10/site-packages/ttax/base_class.py:192, in TT.__getitem__(self, slice_spec)
    187     new_tt_cores.append(sliced_core)
    189 if remainder is not None:
    190   # The reminder obtained from collapsing the last cores.
    191   new_tt_cores[-1] = jnp.einsum('...aib,...bd->...aid',
--> 192                                 new_tt_cores[-1], remainder)
    193   remainder = None
    194 return TT(new_tt_cores)

IndexError: list index out of range
>>> a[0, 0, :]
TT(tt_cores=[DeviceArray([[[0.2888445 ],
                           [0.12369099]]], dtype=float32)])