Merge pull request #3068 from chiamp:api_ref

PiperOrigin-RevId: 529234087

Author: Flax Authors

docs/api_reference/flax.config.rst

@@ -2,6 +2,6 @@
 flax.config package
 ====================
 
-Flax has the following global config options.
-
-.. automodule:: flax.config
+.. automodule:: flax.configurations
+    :members:
+    :exclude-members: use_regular_dict

docs/api_reference/index.rst

@@ -4,13 +4,13 @@ API Reference
 .. toctree::
    :maxdepth: 3
 
+   flax.config
+   flax.core.frozen_dict
+   flax.errors
+   flax.jax_utils
    flax.linen
    flax.serialization
-   flax.core.frozen_dict
    flax.struct
-   flax.jax_utils
    flax.traceback_util
-   flax.traverse_util
    flax.training
-   flax.config
-   flax.errors
+   flax.traverse_util

flax/configurations.py

@@ -14,10 +14,10 @@
 
 """Global configuration options for Flax.
 
-Now a wrapper over jax.config, in which all config vars have a 'flax_' prefix.
+Now a wrapper over jax.config, in which all config vars have a 'flax\_' prefix.
 
 To modify a config value on run time, call:
-`flax.config.update('flax_<config_name>', <value>)`
+``flax.config.update('flax_<config_name>', <value>)``
 
 """