Using `_is_jax_data` for tree flattening results in incompatibility with some `tree_map` operations

In certain cases, it seems like using `_is_jax_data` as a criterion for flattening trees can lead to structural incompatibilities, which can then result in errors when mapping over trees derived from `distrax` distributions.

To elaborate: let's consider that we have a model represented as a PyTree of parameters and metadata, and that this model contains a `distrax` distribution (or more generally `Jittable`) as a child node. We now wish to perform some selective update or partition operation on our model tree — for instance, to separate the tree into `DeviceArray` leaves and non-`DeviceArray` leaves. To do this, we will first perform a `tree_map` on our existing tree, mapping leaves that match the selection criterion to `True` and leaves that don't match to `False`. We will then use this mapped “_mask_” tree to specify the leaves to set to `None` on either side of the partition.

Unfortunately, this is where we hit a snag. Since our mask tree now contains boolean values in place of `DeviceArray`s, `_is_jax_data` will return False for our mask tree where it returned True for the original tree, and the `children` field could be left empty for the mask tree. Because the flattened distribution and mask trees do not thereafter share the same structure, we cannot use the mask tree as needed to create our partition. (Side note: Even if we didn't create a mask tree for our partition, we'd still end up with `None` on the side of the partition without `DeviceArray`s, ultimately resulting in the same structural incompatibility if we later wish to undo the partition.) I'm not actually sure whether the data-based flattening switch is the only cause here, but wanted to share my observations.

Here is a minimal reproducible example demonstrating the issue:
```python
import jax, distrax

tree = distrax.Normal(0, 1)
mask = jax.tree_util.tree_map(lambda _: True, tree)

jax.tree_util.tree_map(
lambda l, r: l,
tree,
mask
)
```
Results in:
```
ValueError: Mismatch custom node data:
([None, None], [True, True], PyTreeDef({'_loc': *, '_scale': *})) !=
([True, True], [False, False], PyTreeDef({'_loc': *, '_scale': *}));
value: .
```
As a result of this design choice, `distrax` distributions are not currently compatible with `equinox`’s [filter transforms](https://docs.kidger.site/equinox/api/filtering/filtered-transformations/), like `eqx.filter_jit`. This doesn't actually matter much for my use case — I can mark any model fields that are `distrax.Distribution`s as static without recompiling since the instance doesn't change — but it is possible there are other use cases where this could make a difference.

## Details

JAX v0.3.16
distrax v0.1.2 (nightly from c013670ad5690f1d07bfa7fcd2a1091f54b901ef)
Running on CPU