Files

.github
benchmarks
build
cloud_tpu_colabs
docs
- _static
- _templates
- _tutorials
- cuda_custom_call
- debugging
- export
- ffi
- gpu_ops
- jep
- notebooks
- pallas
- sphinxext
- Custom_Operation_for_GPUs.md
- Custom_Operation_for_GPUs.py
- README.md
- advanced-autodiff.md
- advanced_guide.rst
- aot.md
- api_compatibility.md
- async_dispatch.rst
- autodidax.ipynb
- autodidax.md
- autodidax.py
- automatic-differentiation.md
- automatic-vectorization.md
- beginner_guide.rst
- build_custom_gpu.sh
- building_on_jax.md
- changelog.md
- concurrency.rst
- conf.py
- contributing.md
- contributor_guide.rst
- debugging.md
- deprecation.md
- developer.md
- device_memory_profiling.md
- distributed_data_loading.md
- errors.rst
- extensions.rst
- external-callbacks.md
- faq.rst
- ffi.ipynb
- ffi.md
- glossary.rst
- gpu_memory_allocation.rst
- gpu_performance_tips.md
- gradient-checkpointing.md
- index.rst
- installation.md
- investigating_a_regression.md
- jax-primitives.md
- jax.debug.rst
- jax.distributed.rst
- jax.dlpack.rst
- jax.dtypes.rst
- jax.example_libraries.optimizers.rst
- jax.example_libraries.rst
- jax.example_libraries.stax.rst
- jax.experimental.array_api.rst
- jax.experimental.checkify.rst
- jax.experimental.compilation_cache.rst
- jax.experimental.custom_partitioning.rst
- jax.experimental.host_callback.rst
- jax.experimental.jet.rst
- jax.experimental.key_reuse.rst
- jax.experimental.mesh_utils.rst
- jax.experimental.multihost_utils.rst
- jax.experimental.pallas.rst
- jax.experimental.pjit.rst
- jax.experimental.rst
- jax.experimental.serialize_executable.rst
- jax.experimental.shard_map.rst
- jax.experimental.sparse.rst
- jax.export.rst
- jax.extend.ffi.rst
- jax.extend.linear_util.rst
- jax.extend.mlir.rst
- jax.extend.random.rst
- jax.extend.rst
- jax.flatten_util.rst
- jax.image.rst
- jax.lax.rst
- jax.lib.rst
- jax.nn.initializers.rst
- jax.nn.rst
- jax.numpy.rst
- jax.ops.rst
- jax.profiler.rst
- jax.random.rst
- jax.rst
- jax.scipy.rst
- jax.sharding.rst
- jax.stages.rst
- jax.tree.rst
- jax.tree_util.rst
- jax.typing.rst
- jax_array_migration.md
- jax_internal_api.rst
- jaxpr.md
- jit-compilation.md
- key-concepts.md
- multi_process.md
- notes.rst
- persistent_compilation_cache.md
- profiling.md
- pytrees.md
- quickstart.md
- random-numbers.md
- rank_promotion_warning.rst
- requirements.txt
- sharded-computation.ipynb
- sharded-computation.md
- stateful-computations.md
- transfer_guard.rst
- tutorials.rst
- type_promotion.rst
- user_guides.rst
- working-with-pytrees.md
- xla_flags.md
examples
images
jax
jax_plugins
jaxlib
tests
third_party
.bazelrc
.bazelversion
.editorconfig
.gitignore
.pre-commit-config.yaml
.readthedocs.yml
AUTHORS
CHANGELOG.md
CITATION.bib
CONTRIBUTING.md
LICENSE
README.md
WORKSPACE
conftest.py
platform_mappings
pyproject.toml
setup.py

jax.example_libraries.stax.rst

jakevdp

DOC: make API doc titles more uniform

Jan 18, 2023

81e627d · Jan 18, 2023

History

`jax.example_libraries.stax` module

.. automodule:: jax.example_libraries.stax
    :members:
    :undoc-members:
    :show-inheritance:

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Files

jax.example_libraries.stax.rst

jax.example_libraries.stax.rst

`jax.example_libraries.stax` module

Files

jax.example_libraries.stax.rst

Latest commit

History

jax.example_libraries.stax.rst

File metadata and controls

jax.example_libraries.stax module

`jax.example_libraries.stax` module