`TracerArrayConversionError` when indexing into a distribution

I sometimes get an error when indexing into a `Distribution` object via square brackets.
For example, here's a snippet that breaks:

```python
from jax import jit
import jax.numpy as jnp
import distrax

def foo(logits):
d = distrax.Categorical(logits=logits)
mask = jnp.arange(logits.shape[0])
return d[mask].entropy()

logits = jnp.zeros((3, 3, 3))
foo(logits) # fine
jit(foo)(logits) # raises
```

Specifically I get the following traceback:

```
Traceback (most recent call last):
File "", line 1, in
File "", line 3, in foo
File ".venv/lib64/python3.11/site-packages/distrax/_src/distributions/categorical.py", line 156, in __getitem__
index = distribution.to_batch_shape_index(self.batch_shape, index)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File ".venv/lib64/python3.11/site-packages/distrax/_src/distributions/distribution.py", line 354, in to_batch_shape_index
new_index = [x[index] for x in np.indices(batch_shape)]
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File ".venv/lib64/python3.11/site-packages/distrax/_src/distributions/distribution.py", line 354, in
new_index = [x[index] for x in np.indices(batch_shape)]
~^^^^^^^
jax.errors.TracerArrayConversionError: The numpy.ndarray conversion method __array__() was called on traced array with shape bool[3,3]
The error occurred while tracing the function foo at :1 for jit. This value became a tracer due to JAX operations on these lines:

operation a:i32[] = convert_element_type[new_dtype=int32 weak_type=True] b
from line :3:11 (foo)
See https://jax.readthedocs.io/en/latest/errors.html#jax.errors.TracerArrayConversionError
--------------------
For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.
```

It seems the error is in [this line](https://github.com/google-deepmind/distrax/blob/7fc5bd7efff4a7144d175199159f115c3e68a3cf/distrax/_src/distributions/distribution.py#L354).
The error goes away if I change `np.indices` to `jnp.indices` in that line.
Would that be an available solution?