.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/examples.rst
docs/examples_community_examples.rst
docs/examples_core_examples.rst
docs/examples_google_research_examples.rst
docs/examples_repositories_that_use_flax.rst
docs/flax.png
docs/getting_started.ipynb
docs/getting_started.md
docs/glossary.rst
docs/index.rst
docs/mission.md
docs/overview.md
docs/philosophy.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.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/flip/0000-template.md
docs/flip/1009-optimizer-api.md
docs/flip/1777-default-dtype.md
docs/flip/2434-general-metadata.md
docs/flip/2974-kw-only-dataclasses.md
docs/flip/README.md
docs/guides/arguments.md
docs/guides/batch_norm.rst
docs/guides/convert_pytorch_to_flax.rst
docs/guides/dropout.rst
docs/guides/ensembling.rst
docs/guides/extracting_intermediates.rst
docs/guides/flax_basics.ipynb
docs/guides/flax_basics.md
docs/guides/flax_on_pjit.ipynb
docs/guides/flax_on_pjit.md
docs/guides/full_eval.rst
docs/guides/haiku_migration_guide.rst
docs/guides/index.rst
docs/guides/index_converting_and_upgrading.rst
docs/guides/index_data_preprocessing.rst
docs/guides/index_flax_fundamentals.rst
docs/guides/index_model_inspection.rst
docs/guides/index_parallel_training.rst
docs/guides/index_training_techniques.rst
docs/guides/jax_for_the_impatient.ipynb
docs/guides/jax_for_the_impatient.md
docs/guides/linen_upgrade_guide.rst
docs/guides/lr_schedule.rst
docs/guides/model_surgery.ipynb
docs/guides/model_surgery.md
docs/guides/optax_update_guide.rst
docs/guides/orbax_upgrade_guide.rst
docs/guides/regular_dict_upgrade_guide.rst
docs/guides/rnncell_upgrade_guide.rst
docs/guides/setup_or_nncompact.rst
docs/guides/state_params.rst
docs/guides/transfer_learning.ipynb
docs/guides/transfer_learning.md
docs/guides/use_checkpointing.ipynb
docs/guides/use_checkpointing.md
docs/notebooks/flax_sharp_bits.ipynb
docs/notebooks/flax_sharp_bits.md
docs/notebooks/full_eval.ipynb
docs/notebooks/full_eval.md
docs/notebooks/index.rst
docs/notebooks/linen_intro.ipynb
docs/notebooks/linen_intro.md
docs/notebooks/optax_update_guide.ipynb
docs/notebooks/optax_update_guide.md
docs/notebooks/orbax_upgrade_guide.ipynb
docs/notebooks/orbax_upgrade_guide.md
docs/notebooks/state_params.ipynb
docs/notebooks/state_params.md
examples/README.md
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/linen_design_test/tied_autoencoder.py
examples/linen_design_test/weight_std.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/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/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/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/linen/README.md
flax/linen/__init__.py
flax/linen/activation.py
flax/linen/attention.py
flax/linen/combinators.py
flax/linen/dtypes.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/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_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
tests/linen/toplevel_test.py