Fix docs build.

PiperOrigin-RevId: 548953086

docs/api.rst

@@ -223,7 +223,7 @@ Params
 .. autoclass:: Params
 
 MutableParams
-~~~~~~
+~~~~~~~~~~~~~
 
 .. autoclass:: MutableParams
 
@@ -233,7 +233,7 @@ State
 .. autoclass:: State
 
 MutableState
-~~~~~~
+~~~~~~~~~~~~
 
 .. autoclass:: MutableState
 
@@ -941,6 +941,7 @@ Control Flow
 
     cond
     fori_loop
+    map
     scan
     switch
     while_loop
@@ -955,6 +956,11 @@ fori_loop
 
 .. autofunction:: fori_loop
 
+map
+~~~~~~~~~
+
+.. autofunction:: map
+
 scan
 ~~~~
 
@@ -1197,12 +1203,12 @@ get_current_state
 .. autofunction:: get_current_state
 
 maybe_get_rng_sequence_state
-~~~~~~~~~~~~~~~~~~
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
 .. autofunction:: maybe_get_rng_sequence_state
 
 replace_rng_sequence_state
-~~~~~~~~~~~~~~~~~~
+~~~~~~~~~~~~~~~~~~~~~~~~~~
 
 .. autofunction:: replace_rng_sequence_state
 
@@ -1213,6 +1219,7 @@ Optimizations
 
     optimize_rng_use
     layer_stack
+    LayerStackTransparencyMapping
     module_auto_repr
     fast_eval_shape
     rng_reserve_size
@@ -1227,6 +1234,11 @@ layer_stack
 
 .. autofunction:: layer_stack
 
+LayerStackTransparencyMapping
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: LayerStackTransparencyMapping
+
 module_auto_repr
 ~~~~~~~~~~~~~~~~
 
@@ -1238,7 +1250,7 @@ fast_eval_shape
 .. autofunction:: fast_eval_shape
 
 rng_reserve_size
-~~~~~~~~~~~~~~~
+~~~~~~~~~~~~~~~~
 
 .. autofunction:: rng_reserve_size
 

docs/conf.py

@@ -231,3 +231,6 @@ def linkcode_resolve(domain, info):
     __ https://github.com/deepmind/dm-haiku/blob/
         {{ env.config.release }}/{{ docname }}
 """
+
+# -- bibtex configuration ---------------------------------------------------
+bibtex_bibfiles = ['references.bib']

docs/requirements.txt

@@ -1,4 +1,4 @@
-absl-py==0.15.0
+absl-py==1.4.0
 alabaster==0.7.12
 attrs==21.2.0
 Babel==2.9.1
@@ -18,12 +18,11 @@ importlib-resources==5.9.0
 ipykernel==5.3.4
 ipython==7.16.1
 ipython-genutils==0.2.0
-jax==0.3.16
-jaxlib==0.3.15
+jax==0.4.13
+jaxlib==0.4.13
 jedi==0.18.0
 Jinja2==2.11.3
-jmp==0.0.2
-jq==1.1.1
+jmp==0.0.4
 jsonschema==4.1.2
 jupyter-client==7.0.6
 jupyter-core==4.9.1
@@ -38,7 +37,7 @@ nbconvert==6.2.0
 nbformat==5.1.3
 nbsphinx==0.8.9
 nest-asyncio==1.5.1
-numpy==1.21.3
+numpy==1.24.0
 opt-einsum==3.3.0
 oset==0.1.3
 packaging==21.2
@@ -60,17 +59,17 @@ pyrsistent==0.18.0
 python-dateutil==2.8.2
 pytz==2021.3
 PyYAML==6.0
-pyzmq==22.3.0
+pyzmq==23.0.0
 requests==2.26.0
-scipy==1.7.1
+scipy==1.11.1
 seaborn==0.11.1
 six==1.16.0
 snowballstemmer==2.1.0
 sphinx==4.5.0
 sphinx-autodoc-typehints==1.11.1
 sphinx-book-theme==0.3.3
 sphinxcontrib-applehelp==1.0.2
-sphinxcontrib-bibtex==1.0.0
+sphinxcontrib-bibtex==2.5.0
 sphinxcontrib-devhelp==1.0.2
 sphinxcontrib-htmlhelp==2.0.0
 sphinxcontrib-jsmath==1.0.1

haiku/_src/lift.py

@@ -421,9 +421,9 @@ def lift_with_state(
   See also:
     :func:`~haiku.lift`: Register parameters with an outer transform.
     :func:`transparent_lift`: Register parameters with an outer transform
-      without a namespace.
+    without a namespace.
     :func:`transparent_lift_with_state`: Register parameters and state with an
-      outer transform without a namespace.
+    outer transform without a namespace.
   """
   base.assert_context("experimental.lift_with_state")
   params_and_state_fn = _to_callable(

haiku/_src/stateful.py

@@ -541,7 +541,7 @@ def cond(*args, **kwargs):
 
 @with_output_structure_hint
 def switch(index, branches, *operands):
-  """Equivalent to :func:`jax.lax.switch` but with Haiku state passed in/out.
+  r"""Equivalent to :func:`jax.lax.switch` but with Haiku state passed in/out.
 
   Note that creating parameters inside a switch branch is not supported, as such
   at init time we recommend you unconditionally evaluate all branches of your
@@ -564,7 +564,7 @@ def switch(index, branches, *operands):
     operands: Operands (A) input to whichever branch is applied.
 
   Returns:
-    Value (B) of branch(*operands) for the branch that was selected based on
+    Value (B) of branch(\*operands) for the branch that was selected based on
     index.
   """
   if not base.inside_transform():