Unexpected behaviour when passing a `Distribution` to a function with `donate_argnums`

I'm having trouble passing a Distribution object to a function when its argument is annotated with `donate_argnums`:

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

x = Categorical(logits=jnp.zeros(4))

def foo(x):
return x.logits

f = jax.jit(foo, donate_argnums=0).trace(x).lower().compile()
f(x)
```

Here's the traceback:

```none
Traceback (most recent call last):
File "foo.py", line 13, in
f(x)
File ".venv/lib64/python3.11/site-packages/jax/_src/stages.py", line 849, in __call__
return self._call(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^
File ".venv/lib64/python3.11/site-packages/jax/_src/interpreters/pxla.py", line 3277, in aot_cache_miss
outs, out_flat, args_flat = stages.Compiled.call(params, *args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File ".venv/lib64/python3.11/site-packages/jax/_src/stages.py", line 813, in call
raise TypeError('\n'.join(msg))
TypeError: Function compiled with input pytree does not match the input pytree it was called with. There are 1 mismatches, including:
* at args[0], seen with pytree metadata ([, ArgInfo(_aval=ShapedArray(float32[4]), donated=True)], [False, False], PyTreeDef({'_dtype': *, '_logits': *, '_probs': None})) but now given with pytree metadata ([, ], [False, False], PyTreeDef({'_dtype': *, '_logits': *, '_probs': None})), so the pytree node metadata does not match
```

The code works fine if I explicitly flatten the `Distribution` and unflatten inside jit:

```python
def foo(leaves, treedef):
x = jax.tree.unflatten(treedef, leaves)
return x.logits

leaves, treedef = jax.tree.flatten(x)
f = (
jax.jit(foo, donate_argnums=0, static_argnums=1)
.trace(leaves, treedef)
.lower()
.compile()
)
f(leaves)
```

But this is quite cumbersome. I think the correct approach is to also mark `jax.stages.ArgInfo` as a form of "Jax data" inside [`distrax._src.utils.jittable._is_jax_data`](https://github.com/google-deepmind/distrax/blob/d9010575b648d800f2dcb7a9ecfb13976e6824df/distrax/_src/utils/jittable.py#L53)

Indeed, after making this change, the original code works as expected.