Merge pull request #733 from fabianp:doctest

PiperOrigin-RevId: 601143290

docs/conf.py

@@ -145,6 +145,13 @@ def new_process_docstring(app, what, name, obj, options, lines):
     'sphinxcontrib.collections'
 ]
 
+# so we don't have to do the canonical imports on every doctest
+doctest_global_setup = '''
+import optax
+import jax
+import jax.numpy as jnp
+'''
+
 # Add any paths that contain templates here, relative to this directory.
 templates_path = ['_templates']
 

optax/tree_utils/_state_utils.py

@@ -51,16 +51,16 @@ def tree_map_params(
   `transform_non_params` can be used to replace any remaining fields as
   required, in this case, we replace those fields by None.
 
-  >>> params, specs = ...  # Trees with the same shape
+  >>> params, specs = jnp.array(0.), jnp.array(0.)  # Trees with the same shape
+  >>> opt = optax.sgd(1e-3)
   >>> state = opt.init(params)
-  >>>
   >>> opt_specs = optax.tree_map_params(
-  >>>     opt,
-  >>>     lambda _, spec: spec,
-  >>>     state,
-  >>>     specs,
-  >>>     transform_non_params=lambda _: None,
-  >>> )
+  ...     opt,
+  ...     lambda _, spec: spec,
+  ...     state,
+  ...     specs,
+  ...     transform_non_params=lambda _: None,
+  ...     )
 
   Args:
     initable: A callable taking parameters and returning an optimizer state, or

optax/tree_utils/_tree_math.py

@@ -119,10 +119,10 @@ def tree_vdot(tree_x: Any, tree_y: Any) -> chex.Numeric:
     inner product between ``tree_x`` and ``tree_y``, a scalar value.
 
   >>> optax.tree_utils.tree_vdot(
-  >>>   {a: jnp.array([1, 2]), b: jnp.array([1, 2])},
-  >>>   {a: jnp.array([-1, -1]), b: jnp.array([1, 1])},
-  >>> )
-  0.0
+  ...   {'a': jnp.array([1, 2]), 'b': jnp.array([1, 2])},
+  ...   {'a': jnp.array([-1, -1]), 'b': jnp.array([1, 1])},
+  ... )
+  Array(0, dtype=int32)
 
   Implementation detail: we upcast the values to the highest precision to avoid
   numerical issues.

test.sh

@@ -76,6 +76,8 @@ cd ..
 # Build Sphinx docs.
 pip install -e ".[docs]"
 cd docs && make html
+# run doctests
+make doctest
 cd ..
 
 set +u