`flax.jax_utils` not discoverable from the NNX docs.

[ibbyml]

Description

Hi! I’m wondering if this is intended: flax.jax_utils is not discoverable via the new NNX docs, even though it’s still part of the flax top level module and actively maintained.

On the Linen docs, flax.jax_utils shows up as expected:

On the NNX docs, searching for jax_utils returns no results:

From a user perspective, this makes it look like flax.jax_utils has been removed or deprecated in favor of a new NNX specific helper, even though that doesn’t seem to be the case.

Is this just a docs/search configuration issue (e.g. NNX docs not indexing flax.jax_utils), or is flax.jax_utils meant to be considered Linen only going forward?

Thanks.

[vfdev-5]

Thanks for the question @ibbyml. This module contains a code using now deprecated jax API. Test coverage of the module is not great neither. I'll ask the team whether it makes sense to expose it in the NNX docs or it is hidden on purpose.

[vfdev-5]

@ibbyml yes, we should expose jax_utils module the NNX docs as well. If you would like to make a contribution it would be helpful!

[ibbyml]

@ibbyml yes, we should expose jax_utils module the NNX docs as well. If you would like to make a contribution it would be helpful!

Great. I'll open a pull request.