.git-blame-ignore-revs
.gitignore
.pre-commit-config.yaml
.readthedocs.yml
AUTHORS
CHANGELOG.md
LICENSE
README.md
contributing.md
pylintrc
pyproject.toml
.github/pull_request_template.md
.github/ISSUE_TEMPLATE/bug_report.md
.github/analytics/README.md
.github/analytics/get_repo_metrics.py
.github/analytics/issue_activity_since_date.gql
.github/analytics/pr_data_query.gql
.github/analytics/requirements.txt
.github/workflows/build.yml
.github/workflows/pythonpublish.yml
dev/README.md
dev/update_requirements.py
dev/.devcontainer/Dockerfile
dev/.devcontainer/devcontainer.json
docs/.gitignore
docs/Makefile
docs/README.md
docs/conf.py
docs/conf_sphinx_patch.py
docs/contributing.md
docs/faq.rst
docs/flax.png
docs/glossary.rst
docs/index.rst
docs/philosophy.md
docs/quick_start.ipynb
docs/quick_start.md
docs/requirements.txt
docs/_ext/codediff.py
docs/_ext/codediff_test.py
docs/_ext/flax_module.py
docs/_static/css/flax_theme.css
docs/_templates/autosummary/flax_module.rst
docs/api_reference/flax.config.rst
docs/api_reference/flax.core.frozen_dict.rst
docs/api_reference/flax.cursor.rst
docs/api_reference/flax.errors.rst
docs/api_reference/flax.jax_utils.rst
docs/api_reference/flax.serialization.rst
docs/api_reference/flax.struct.rst
docs/api_reference/flax.traceback_util.rst
docs/api_reference/flax.training.rst
docs/api_reference/flax.traverse_util.rst
docs/api_reference/index.rst
docs/api_reference/flax.linen/activation_functions.rst
docs/api_reference/flax.linen/decorators.rst
docs/api_reference/flax.linen/index.rst
docs/api_reference/flax.linen/init_apply.rst
docs/api_reference/flax.linen/initializers.rst
docs/api_reference/flax.linen/inspection.rst
docs/api_reference/flax.linen/layers.rst
docs/api_reference/flax.linen/module.rst
docs/api_reference/flax.linen/profiling.rst
docs/api_reference/flax.linen/spmd.rst
docs/api_reference/flax.linen/transformations.rst
docs/api_reference/flax.linen/variable.rst
docs/developer_notes/index.rst
docs/developer_notes/lift.md
docs/developer_notes/module_lifecycle.rst
docs/examples/community_examples.rst
docs/examples/core_examples.rst
docs/examples/google_research_examples.rst
docs/examples/index.rst
docs/examples/repositories_that_use_flax.rst
docs/experimental/index.rst
docs/flip/0000-template.md
docs/flip/1009-optimizer-api.md
docs/flip/1777-default-dtype.md
docs/flip/2396-rnn.md
docs/flip/2434-general-metadata.md
docs/flip/2974-kw-only-dataclasses.md
docs/flip/README.md
docs/guides/flax_sharp_bits.ipynb
docs/guides/flax_sharp_bits.md
docs/guides/index.rst
docs/guides/converting_and_upgrading/convert_pytorch_to_flax.rst
docs/guides/converting_and_upgrading/haiku_migration_guide.rst
docs/guides/converting_and_upgrading/index.rst
docs/guides/converting_and_upgrading/linen_upgrade_guide.rst
docs/guides/converting_and_upgrading/optax_update_guide.rst
docs/guides/converting_and_upgrading/orbax_upgrade_guide.rst
docs/guides/converting_and_upgrading/regular_dict_upgrade_guide.rst
docs/guides/converting_and_upgrading/rnncell_upgrade_guide.rst
docs/guides/data_preprocessing/full_eval.rst
docs/guides/data_preprocessing/index.rst
docs/guides/data_preprocessing/loading_datasets.ipynb
docs/guides/data_preprocessing/loading_datasets.md
docs/guides/flax_fundamentals/arguments.md
docs/guides/flax_fundamentals/flax_basics.ipynb
docs/guides/flax_fundamentals/flax_basics.md
docs/guides/flax_fundamentals/index.rst
docs/guides/flax_fundamentals/setup_or_nncompact.rst
docs/guides/flax_fundamentals/state_params.rst
docs/guides/model_inspection/extracting_intermediates.rst
docs/guides/model_inspection/index.rst
docs/guides/model_inspection/model_surgery.ipynb
docs/guides/model_inspection/model_surgery.md
docs/guides/parallel_training/ensembling.rst
docs/guides/parallel_training/flax_on_pjit.ipynb
docs/guides/parallel_training/flax_on_pjit.md
docs/guides/parallel_training/index.rst
docs/guides/training_techniques/batch_norm.rst
docs/guides/training_techniques/dropout.rst
docs/guides/training_techniques/index.rst
docs/guides/training_techniques/lr_schedule.rst
docs/guides/training_techniques/transfer_learning.ipynb
docs/guides/training_techniques/transfer_learning.md
docs/guides/training_techniques/use_checkpointing.ipynb
docs/guides/training_techniques/use_checkpointing.md
examples/README.md
examples/__init__.py
examples/cloud/README.md
examples/cloud/launch_gce.py
examples/cloud/startup_script.sh
examples/imagenet/README.md
examples/imagenet/imagenet.ipynb
examples/imagenet/imagenet_benchmark.py
examples/imagenet/imagenet_fake_data_benchmark.py
examples/imagenet/input_pipeline.py
examples/imagenet/main.py
examples/imagenet/models.py
examples/imagenet/models_test.py
examples/imagenet/requirements.txt
examples/imagenet/train.py
examples/imagenet/train_test.py
examples/imagenet/configs/default.py
examples/imagenet/configs/fake_data_benchmark.py
examples/imagenet/configs/tpu.py
examples/imagenet/configs/v100_x8.py
examples/imagenet/configs/v100_x8_mixed_precision.py
examples/linen_design_test/attention_simple.py
examples/linen_design_test/autoencoder.py
examples/linen_design_test/dense.py
examples/linen_design_test/linear_regression.py
examples/linen_design_test/mlp_explicit.py
examples/linen_design_test/mlp_inline.py
examples/linen_design_test/mlp_lazy.py
examples/lm1b/README.md
examples/lm1b/input_pipeline.py
examples/lm1b/input_pipeline_test.py
examples/lm1b/main.py
examples/lm1b/models.py
examples/lm1b/requirements.txt
examples/lm1b/temperature_sampler.py
examples/lm1b/temperature_sampler_test.py
examples/lm1b/tokenizer.py
examples/lm1b/train.py
examples/lm1b/train_test.py
examples/lm1b/utils.py
examples/lm1b/configs/default.py
examples/mnist/README.md
examples/mnist/main.py
examples/mnist/mnist.ipynb
examples/mnist/mnist_benchmark.py
examples/mnist/requirements.txt
examples/mnist/train.py
examples/mnist/train_test.py
examples/mnist/configs/default.py
examples/nlp_seq/README.md
examples/nlp_seq/input_pipeline.py
examples/nlp_seq/input_pipeline_test.py
examples/nlp_seq/models.py
examples/nlp_seq/requirements.txt
examples/nlp_seq/train.py
examples/ogbg_molpcba/README.md
examples/ogbg_molpcba/input_pipeline.py
examples/ogbg_molpcba/input_pipeline_test.py
examples/ogbg_molpcba/main.py
examples/ogbg_molpcba/models.py
examples/ogbg_molpcba/models_test.py
examples/ogbg_molpcba/ogbg_molpcba.ipynb
examples/ogbg_molpcba/ogbg_molpcba_benchmark.py
examples/ogbg_molpcba/requirements.txt
examples/ogbg_molpcba/train.py
examples/ogbg_molpcba/train_test.py
examples/ogbg_molpcba/configs/default.py
examples/ogbg_molpcba/configs/default_graph_net.py
examples/ogbg_molpcba/configs/hparam_sweep.py
examples/ogbg_molpcba/configs/test.py
examples/ppo/README.md
examples/ppo/agent.py
examples/ppo/env_utils.py
examples/ppo/models.py
examples/ppo/ppo_lib.py
examples/ppo/ppo_lib_test.py
examples/ppo/ppo_main.py
examples/ppo/requirements.txt
examples/ppo/seed_rl_atari_preprocessing.py
examples/ppo/test_episodes.py
examples/ppo/configs/default.py
examples/seq2seq/README.md
examples/seq2seq/input_pipeline.py
examples/seq2seq/models.py
examples/seq2seq/requirements.txt
examples/seq2seq/seq2seq.ipynb
examples/seq2seq/train.py
examples/seq2seq/train_test.py
examples/sst2/README.md
examples/sst2/build_vocabulary.py
examples/sst2/input_pipeline.py
examples/sst2/input_pipeline_test.py
examples/sst2/main.py
examples/sst2/models.py
examples/sst2/models_test.py
examples/sst2/requirements.txt
examples/sst2/sst2.ipynb
examples/sst2/train.py
examples/sst2/train_test.py
examples/sst2/vocab.txt
examples/sst2/vocabulary.py
examples/sst2/configs/default.py
examples/vae/README.md
examples/vae/input_pipeline.py
examples/vae/main.py
examples/vae/models.py
examples/vae/reconstruction.png
examples/vae/requirements.txt
examples/vae/sample.png
examples/vae/train.py
examples/vae/utils.py
examples/vae/configs/default.py
examples/vae/results/.gitignore
examples/wmt/README.md
examples/wmt/bleu.py
examples/wmt/decode.py
examples/wmt/input_pipeline.py
examples/wmt/input_pipeline_test.py
examples/wmt/main.py
examples/wmt/models.py
examples/wmt/requirements.txt
examples/wmt/tokenizer.py
examples/wmt/train.py
examples/wmt/train_test.py
examples/wmt/configs/default.py
flax/__init__.py
flax/configurations.py
flax/cursor.py
flax/errors.py
flax/ids.py
flax/io.py
flax/jax_utils.py
flax/py.typed
flax/serialization.py
flax/struct.py
flax/traceback_util.py
flax/traverse_util.py
flax/version.py
flax.egg-info/PKG-INFO
flax.egg-info/SOURCES.txt
flax.egg-info/dependency_links.txt
flax.egg-info/requires.txt
flax.egg-info/top_level.txt
flax/core/__init__.py
flax/core/axes_scan.py
flax/core/flax_functional_engine.ipynb
flax/core/frozen_dict.py
flax/core/lift.py
flax/core/meta.py
flax/core/partial_eval.py
flax/core/scope.py
flax/core/tracers.py
flax/core/variables.py
flax/core/nn/__init__.py
flax/core/nn/attention.py
flax/core/nn/linear.py
flax/core/nn/normalization.py
flax/core/nn/stochastic.py
flax/experimental/__init__.py
flax/experimental/nnx/.gitignore
flax/experimental/nnx/README.md
flax/experimental/nnx/__init__.py
flax/experimental/nnx/docs/blog.md
flax/experimental/nnx/docs/quick_start.ipynb
flax/experimental/nnx/docs/tiny_nnx.ipynb
flax/experimental/nnx/docs/why.ipynb
flax/experimental/nnx/docs/why.md
flax/experimental/nnx/docs/images/stateful-transforms.png
flax/experimental/nnx/examples/lm1b/README.md
flax/experimental/nnx/examples/lm1b/input_pipeline.py
flax/experimental/nnx/examples/lm1b/input_pipeline_test.py
flax/experimental/nnx/examples/lm1b/main.py
flax/experimental/nnx/examples/lm1b/models.py
flax/experimental/nnx/examples/lm1b/models_test.py
flax/experimental/nnx/examples/lm1b/requirements.txt
flax/experimental/nnx/examples/lm1b/temperature_sampler.py
flax/experimental/nnx/examples/lm1b/temperature_sampler_test.py
flax/experimental/nnx/examples/lm1b/tokenizer.py
flax/experimental/nnx/examples/lm1b/train.py
flax/experimental/nnx/examples/lm1b/train_test.py
flax/experimental/nnx/examples/lm1b/utils.py
flax/experimental/nnx/examples/lm1b/configs/default.py
flax/experimental/nnx/examples/toy_examples/00_demo.ipynb
flax/experimental/nnx/examples/toy_examples/01_functional_api.py
flax/experimental/nnx/examples/toy_examples/02_lifted_transforms.py
flax/experimental/nnx/examples/toy_examples/03_train_state.py
flax/experimental/nnx/examples/toy_examples/05_vae.py
flax/experimental/nnx/examples/toy_examples/06_scan_over_layers.py
flax/experimental/nnx/examples/toy_examples/07_transformer.py
flax/experimental/nnx/examples/toy_examples/08_save_load_checkpoints.py
flax/experimental/nnx/examples/toy_examples/09_parameter_surgery.py
flax/experimental/nnx/examples/toy_examples/10_quantization.py
flax/experimental/nnx/examples/toy_examples/requirements.txt
flax/experimental/nnx/ideas/shape_inference.py
flax/experimental/nnx/nnx/__init__.py
flax/experimental/nnx/nnx/compatibility.py
flax/experimental/nnx/nnx/dataclasses.py
flax/experimental/nnx/nnx/errors.py
flax/experimental/nnx/nnx/filterlib.py
flax/experimental/nnx/nnx/flaglib.py
flax/experimental/nnx/nnx/graph_utils.py
flax/experimental/nnx/nnx/helpers.py
flax/experimental/nnx/nnx/ids.py
flax/experimental/nnx/nnx/module.py
flax/experimental/nnx/nnx/proxy_caller.py
flax/experimental/nnx/nnx/pytreelib.py
flax/experimental/nnx/nnx/reprlib.py
flax/experimental/nnx/nnx/rnglib.py
flax/experimental/nnx/nnx/spmd.py
flax/experimental/nnx/nnx/state.py
flax/experimental/nnx/nnx/tracers.py
flax/experimental/nnx/nnx/transforms.py
flax/experimental/nnx/nnx/variables.py
flax/experimental/nnx/nnx/nn/__init__.py
flax/experimental/nnx/nnx/nn/activations.py
flax/experimental/nnx/nnx/nn/attention.py
flax/experimental/nnx/nnx/nn/dtypes.py
flax/experimental/nnx/nnx/nn/initializers.py
flax/experimental/nnx/nnx/nn/linear.py
flax/experimental/nnx/nnx/nn/normalization.py
flax/experimental/nnx/nnx/nn/stochastic.py
flax/experimental/nnx/scripts/requirements.txt
flax/experimental/nnx/scripts/run-all-examples.bash
flax/experimental/nnx/tests/__init__.py
flax/experimental/nnx/tests/test_compatibility.py
flax/experimental/nnx/tests/test_containers.py
flax/experimental/nnx/tests/test_graph_utils.py
flax/experimental/nnx/tests/test_helpers.py
flax/experimental/nnx/tests/test_ids.py
flax/experimental/nnx/tests/test_integration.py
flax/experimental/nnx/tests/test_module.py
flax/experimental/nnx/tests/test_partitioning.py
flax/experimental/nnx/tests/test_pytree.py
flax/experimental/nnx/tests/test_rngs.py
flax/experimental/nnx/tests/test_spmd.py
flax/experimental/nnx/tests/test_state.py
flax/experimental/nnx/tests/test_transforms.py
flax/experimental/nnx/tests/test_variable.py
flax/experimental/nnx/tests/nn/test_attention.py
flax/experimental/nnx/tests/nn/test_conv.py
flax/experimental/nnx/tests/nn/test_embed.py
flax/experimental/nnx/tests/nn/test_linear.py
flax/linen/README.md
flax/linen/__init__.py
flax/linen/activation.py
flax/linen/attention.py
flax/linen/batch_apply.py
flax/linen/combinators.py
flax/linen/dtypes.py
flax/linen/fp8_ops.py
flax/linen/initializers.py
flax/linen/kw_only_dataclasses.py
flax/linen/linear.py
flax/linen/module.py
flax/linen/normalization.py
flax/linen/partitioning.py
flax/linen/pooling.py
flax/linen/recurrent.py
flax/linen/spmd.py
flax/linen/stochastic.py
flax/linen/summary.py
flax/linen/transforms.py
flax/linen/experimental/layers_with_named_axes.py
flax/metrics/__init__.py
flax/metrics/tensorboard.py
flax/oss/ .git-blame-ignore-revs
flax/testing/__init__.py
flax/testing/benchmark.py
flax/training/__init__.py
flax/training/checkpoints.py
flax/training/common_utils.py
flax/training/dynamic_scale.py
flax/training/early_stopping.py
flax/training/lr_schedule.py
flax/training/orbax_utils.py
flax/training/prefetch_iterator.py
flax/training/train_state.py
images/flax_logo.png
images/flax_logo.svg
images/flax_logo_250px.png
images/flax_logo_500px.png
tests/checkpoints_test.py
tests/colab_tpu_jax_version.ipynb
tests/configurations_test.py
tests/cursor_test.py
tests/download_dataset_metadata.sh
tests/early_stopping_test.py
tests/import_test.ipynb
tests/io_test.py
tests/jax_utils_test.py
tests/run_all_tests.sh
tests/serialization_test.py
tests/struct_test.py
tests/tensorboard_test.py
tests/traceback_util_test.py
tests/traverse_util_test.py
tests/core/core_frozen_dict_test.py
tests/core/core_lift_test.py
tests/core/core_meta_test.py
tests/core/core_scope_test.py
tests/core/design/core_attention_test.py
tests/core/design/core_auto_encoder_test.py
tests/core/design/core_big_resnets_test.py
tests/core/design/core_custom_vjp_test.py
tests/core/design/core_dense_test.py
tests/core/design/core_flow_test.py
tests/core/design/core_resnet_test.py
tests/core/design/core_scan_test.py
tests/core/design/core_tied_autoencoder_test.py
tests/core/design/core_vmap_test.py
tests/core/design/core_weight_std_test.py
tests/linen/initializers_test.py
tests/linen/kw_only_dataclasses_test.py
tests/linen/linen_activation_test.py
tests/linen/linen_attention_test.py
tests/linen/linen_batch_apply_test.py
tests/linen/linen_combinators_test.py
tests/linen/linen_dtypes_test.py
tests/linen/linen_linear_test.py
tests/linen/linen_meta_test.py
tests/linen/linen_module_test.py
tests/linen/linen_recurrent_test.py
tests/linen/linen_test.py
tests/linen/linen_transforms_test.py
tests/linen/partitioning_test.py
tests/linen/summary_test.py