# Helper functions for Equinox filtering

To make use of the full power of JAX, it is highly recommended to learn about [equinox](https://docs.kidger.site/equinox/all-of-equinox/). Using `equinox`, cryoJAX implements its [models as pytrees](https://docs.kidger.site/equinox/all-of-equinox/#1-models-as-pytrees) using `equinox.Module`s. These pytrees can be operated on similarly to any pytree with JAX (e.g. with `jax.tree.map`). Complementary to the `equinox.Module` interface, `equinox` introduces the idea of [*filtering*](https://docs.kidger.site/equinox/all-of-equinox/#2-filtering) in order to separate pytree leaves into different groups, a core task in using JAX (e.g. separating traced and static arguments to `jax.jit`). In particular, this grouping is achieved with the functions [eqx.partition and eqx.combine](https://docs.kidger.site/equinox/api/manipulation/#equinox.partition). This documentation describes utilities in `cryojax` for working with `equinox.partition` and `equinox.combine`.

::: cryojax.jax_util.get_filter_spec
