jax.experimental.stax module .. automodule:: jax.experimental.stax :members: :undoc-members: :show-inheritance: