rocm_jax/docs/jax.example_libraries.stax.rst

``jax.example_libraries.stax`` module
=====================================

.. automodule:: jax.example_libraries.stax
    :members:
    :undoc-members:
    :show-inheritance:
DOC: make API doc titles more uniform 2023-01-18 10:59:42 -08:00			``jax.example_libraries.stax`` module
			`=====================================`
[JAX] move example libraries from `jax.experimental` into `jax.example_libraries` The `jax.experimental.stax` and `jax.experimental.optimizers` modules are standalone examples libraries. By contrast, the remaining modules in `jax.experimental` are experimental features of the JAX core system. This change moves the two example libraries, and the README that describes them, to `jax.example_libraries` to reflect this distinction. PiperOrigin-RevId: 404405186 2021-10-19 17:30:16 -07:00
			`.. automodule:: jax.example_libraries.stax`
			`:members:`
			`:undoc-members:`
			`:show-inheritance:`