.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/flax_publish.yml
.github/workflows/flax_test.yml
.github/workflows/flaxlib_publish.yml
.github/workflows/jax_nightly.yml
benchmarks/README.md
benchmarks/nnx_graph_overhead.py
benchmarks/nnx_mlpmixer_training.py
benchmarks/nnx_simple_training.py
benchmarks/nnx_state_traversal.py
docs/.gitignore
docs/.readthedocs.yaml
docs/Makefile
docs/README.md
docs/conf.py
docs/conf_sphinx_patch.py
docs/faq.rst
docs/flax.png
docs/glossary.rst
docs/index.rst
docs/linen_intro.ipynb
docs/linen_intro.md
docs/quick_start.ipynb
docs/quick_start.md
docs/robots.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.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/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/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/3099-rnnbase-refactor.md
docs/flip/4105-jax-style-nnx-transforms.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/rng_guide.ipynb
docs/guides/flax_fundamentals/rng_guide.md
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/quantization/fp8_basics.ipynb
docs/guides/quantization/fp8_basics.md
docs/guides/quantization/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
docs_nnx/.gitignore
docs_nnx/.readthedocs.yaml
docs_nnx/Makefile
docs_nnx/README.md
docs_nnx/conf.py
docs_nnx/conf_sphinx_patch.py
docs_nnx/contributing.md
docs_nnx/faq.rst
docs_nnx/flax.png
docs_nnx/guides_advanced.rst
docs_nnx/guides_basic.rst
docs_nnx/index.rst
docs_nnx/key_concepts.ipynb
docs_nnx/key_concepts.md
docs_nnx/mnist_tutorial.ipynb
docs_nnx/mnist_tutorial.md
docs_nnx/nnx_basics.ipynb
docs_nnx/nnx_basics.md
docs_nnx/nnx_glossary.rst
docs_nnx/philosophy.md
docs_nnx/robots.txt
docs_nnx/why.rst
docs_nnx/_ext/codediff.py
docs_nnx/_ext/codediff_test.py
docs_nnx/_ext/flax_module.py
docs_nnx/_static/css/flax_theme.css
docs_nnx/_templates/autosummary/flax_module.rst
docs_nnx/api_reference/flax.config.rst
docs_nnx/api_reference/flax.core.frozen_dict.rst
docs_nnx/api_reference/flax.struct.rst
docs_nnx/api_reference/flax.training.rst
docs_nnx/api_reference/flax.traverse_util.rst
docs_nnx/api_reference/index.rst
docs_nnx/api_reference/flax.nnx/bridge.rst
docs_nnx/api_reference/flax.nnx/filterlib.rst
docs_nnx/api_reference/flax.nnx/graph.rst
docs_nnx/api_reference/flax.nnx/helpers.rst
docs_nnx/api_reference/flax.nnx/index.rst
docs_nnx/api_reference/flax.nnx/module.rst
docs_nnx/api_reference/flax.nnx/object.rst
docs_nnx/api_reference/flax.nnx/rnglib.rst
docs_nnx/api_reference/flax.nnx/spmd.rst
docs_nnx/api_reference/flax.nnx/state.rst
docs_nnx/api_reference/flax.nnx/summary.rst
docs_nnx/api_reference/flax.nnx/transforms.rst
docs_nnx/api_reference/flax.nnx/variables.rst
docs_nnx/api_reference/flax.nnx/visualization.rst
docs_nnx/api_reference/flax.nnx/nn/activations.rst
docs_nnx/api_reference/flax.nnx/nn/attention.rst
docs_nnx/api_reference/flax.nnx/nn/dtypes.rst
docs_nnx/api_reference/flax.nnx/nn/index.rst
docs_nnx/api_reference/flax.nnx/nn/initializers.rst
docs_nnx/api_reference/flax.nnx/nn/linear.rst
docs_nnx/api_reference/flax.nnx/nn/lora.rst
docs_nnx/api_reference/flax.nnx/nn/normalization.rst
docs_nnx/api_reference/flax.nnx/nn/recurrent.rst
docs_nnx/api_reference/flax.nnx/nn/stochastic.rst
docs_nnx/api_reference/flax.nnx/training/index.rst
docs_nnx/api_reference/flax.nnx/training/metrics.rst
docs_nnx/api_reference/flax.nnx/training/optimizer.rst
docs_nnx/examples/core_examples.rst
docs_nnx/examples/gemma.ipynb
docs_nnx/examples/gemma.md
docs_nnx/examples/index.rst
docs_nnx/flip/0000-template.md
docs_nnx/flip/1009-optimizer-api.md
docs_nnx/flip/1777-default-dtype.md
docs_nnx/flip/2396-rnn.md
docs_nnx/flip/2434-general-metadata.md
docs_nnx/flip/2974-kw-only-dataclasses.md
docs_nnx/flip/3099-rnnbase-refactor.md
docs_nnx/flip/4105-jax-style-nnx-transforms.md
docs_nnx/flip/4844-var-eager-sharding.md
docs_nnx/flip/README.md
docs_nnx/guides/blog.md
docs_nnx/guides/bridge_guide.ipynb
docs_nnx/guides/bridge_guide.md
docs_nnx/guides/checkpointing.ipynb
docs_nnx/guides/checkpointing.md
docs_nnx/guides/demo.ipynb
docs_nnx/guides/demo.md
docs_nnx/guides/filters_guide.ipynb
docs_nnx/guides/filters_guide.md
docs_nnx/guides/flax_gspmd.ipynb
docs_nnx/guides/flax_gspmd.md
docs_nnx/guides/index.rst
docs_nnx/guides/jax_and_nnx_transforms.rst
docs_nnx/guides/performance.ipynb
docs_nnx/guides/performance.md
docs_nnx/guides/pytree.ipynb
docs_nnx/guides/pytree.md
docs_nnx/guides/randomness.ipynb
docs_nnx/guides/randomness.md
docs_nnx/guides/surgery.ipynb
docs_nnx/guides/surgery.md
docs_nnx/guides/tiny_nnx.ipynb
docs_nnx/guides/transforms.ipynb
docs_nnx/guides/transforms.md
docs_nnx/guides/images/performance-graph.png
docs_nnx/guides/images/stateful-transforms.png
docs_nnx/hijax/hijax.ipynb
docs_nnx/hijax/hijax.md
docs_nnx/hijax/index.rst
docs_nnx/migrating/convert_pytorch_to_flax.rst
docs_nnx/migrating/haiku_to_flax.rst
docs_nnx/migrating/index.rst
docs_nnx/migrating/linen_to_nnx.rst
docs_nnx/migrating/nnx_010_to_nnx_011.rst
examples/README.md
examples/__init__.py
examples/cloud/README.md
examples/cloud/launch_gce.py
examples/cloud/startup_script.sh
examples/gemma/README.md
examples/gemma/helpers.py
examples/gemma/helpers_test.py
examples/gemma/input_pipeline.py
examples/gemma/input_pipeline_test.py
examples/gemma/layers.py
examples/gemma/layers_test.py
examples/gemma/main.py
examples/gemma/modules.py
examples/gemma/modules_test.py
examples/gemma/params.py
examples/gemma/positional_embeddings.py
examples/gemma/positional_embeddings_test.py
examples/gemma/requirements.txt
examples/gemma/sampler.py
examples/gemma/sampler_test.py
examples/gemma/sow_lib.py
examples/gemma/tokenizer.py
examples/gemma/train.py
examples/gemma/transformer.py
examples/gemma/transformer_test.py
examples/gemma/utils.py
examples/gemma/configs/default.py
examples/gemma/configs/gemma3_4b.py
examples/gemma/configs/small.py
examples/gemma/configs/tiny.py
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/lm1b_nnx/README.md
examples/lm1b_nnx/input_pipeline.py
examples/lm1b_nnx/input_pipeline_test.py
examples/lm1b_nnx/main.py
examples/lm1b_nnx/models.py
examples/lm1b_nnx/models_test.py
examples/lm1b_nnx/requirements.txt
examples/lm1b_nnx/temperature_sampler.py
examples/lm1b_nnx/temperature_sampler_test.py
examples/lm1b_nnx/tokenizer.py
examples/lm1b_nnx/train.py
examples/lm1b_nnx/train_test.py
examples/lm1b_nnx/utils.py
examples/lm1b_nnx/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/nnx_toy_examples/01_functional_api.py
examples/nnx_toy_examples/02_lifted_transforms.py
examples/nnx_toy_examples/03_train_state.py
examples/nnx_toy_examples/04_data_parallel_with_jit.py
examples/nnx_toy_examples/05_vae.py
examples/nnx_toy_examples/06_scan_over_layers.py
examples/nnx_toy_examples/07_array_leaves.py
examples/nnx_toy_examples/08_save_load_checkpoints.py
examples/nnx_toy_examples/09_parameter_surgery.py
examples/nnx_toy_examples/10_fsdp_and_optimizer.py
examples/nnx_toy_examples/hijax_basic.py
examples/nnx_toy_examples/hijax_demo.py
examples/nnx_toy_examples/requirements.txt
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/typing.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/spmd.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.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/nnx/README.md
flax/nnx/__init__.py
flax/nnx/extract.py
flax/nnx/filterlib.py
flax/nnx/graph.py
flax/nnx/helpers.py
flax/nnx/ids.py
flax/nnx/module.py
flax/nnx/proxy_caller.py
flax/nnx/pytreelib.py
flax/nnx/reprlib.py
flax/nnx/rnglib.py
flax/nnx/spmd.py
flax/nnx/statelib.py
flax/nnx/summary.py
flax/nnx/tracers.py
flax/nnx/traversals.py
flax/nnx/variablelib.py
flax/nnx/visualization.py
flax/nnx/bridge/__init__.py
flax/nnx/bridge/interop.py
flax/nnx/bridge/module.py
flax/nnx/bridge/variables.py
flax/nnx/bridge/wrappers.py
flax/nnx/nn/__init__.py
flax/nnx/nn/activations.py
flax/nnx/nn/attention.py
flax/nnx/nn/dtypes.py
flax/nnx/nn/initializers.py
flax/nnx/nn/linear.py
flax/nnx/nn/lora.py
flax/nnx/nn/normalization.py
flax/nnx/nn/recurrent.py
flax/nnx/nn/stochastic.py
flax/nnx/scripts/requirements.txt
flax/nnx/scripts/run-all-examples.bash
flax/nnx/training/__init__.py
flax/nnx/training/metrics.py
flax/nnx/training/optimizer.py
flax/nnx/transforms/__init__.py
flax/nnx/transforms/autodiff.py
flax/nnx/transforms/compilation.py
flax/nnx/transforms/general.py
flax/nnx/transforms/iteration.py
flax/nnx/transforms/transforms.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
flaxlib_src/.gitignore
flaxlib_src/CMakeLists.txt
flaxlib_src/Cargo.lock
flaxlib_src/Cargo.toml
flaxlib_src/LICENSE
flaxlib_src/README.md
flaxlib_src/pyproject.toml
flaxlib_src/uv.lock
flaxlib_src/src/lib.cc
flaxlib_src/src/flaxlib/__init__.py
flaxlib_src/src/flaxlib/flaxlib_cpp.pyi
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/flaxlib_test.py
tests/import_test.ipynb
tests/io_test.py
tests/jax_utils_test.py
tests/pickle_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
tests/nnx/__init__.py
tests/nnx/containers_test.py
tests/nnx/filters_test.py
tests/nnx/graph_utils_test.py
tests/nnx/helpers_test.py
tests/nnx/ids_test.py
tests/nnx/integration_test.py
tests/nnx/metrics_test.py
tests/nnx/module_test.py
tests/nnx/mutable_array_test.py
tests/nnx/optimizer_test.py
tests/nnx/partitioning_test.py
tests/nnx/rngs_test.py
tests/nnx/spmd_test.py
tests/nnx/state_test.py
tests/nnx/summary_test.py
tests/nnx/test_traversals.py
tests/nnx/transforms_test.py
tests/nnx/variable_test.py
tests/nnx/bridge/module_test.py
tests/nnx/bridge/wrappers_test.py
tests/nnx/nn/attention_test.py
tests/nnx/nn/conv_test.py
tests/nnx/nn/embed_test.py
tests/nnx/nn/linear_test.py
tests/nnx/nn/lora_test.py
tests/nnx/nn/normalization_test.py
tests/nnx/nn/recurrent_test.py
tests/nnx/nn/stochastic_test.py