docs: add doctest entries

docs/source/index.rst

@@ -16,7 +16,7 @@ OpTree: Optimized PyTree Utilities
     api.rst
 
 .. toctree::
-    :maxdepth: 1
+    :maxdepth: 2
 
     integration.rst
 

optree/integration/jax.py

@@ -117,6 +117,26 @@ def tree_ravel(
 ) -> tuple[Array, Callable[[Array], ArrayTree]]:
     r"""Ravel (flatten) a pytree of arrays down to a 1D array.
 
+    >>> tree = {
+    ...     'layer1': {
+    ...         'weight': jnp.arange(6, dtype=jnp.float32).reshape((2, 3)),
+    ...         'bias': jnp.arange(2, dtype=jnp.float32).reshape((2,)),
+    ...     },
+    ...     'layer2': {
+    ...         'weight': jnp.arange(2, dtype=jnp.float32).reshape((1, 2)),
+    ...         'bias': jnp.arange(1, dtype=jnp.float32).reshape((1,))
+    ...     },
+    ... }
+    >>> flat, unravel_func = tree_ravel(tree)
+    >>> flat
+    Array([0., 1., 0., 1., 2., 3., 4., 5., 0., 0., 1.], dtype=float32)
+    >>> unravel_func(flat)
+    {'layer1': {'weight': Array([[0., 1., 2.],
+                                 [3., 4., 5.]], dtype=float32),
+                'bias': Array([0., 1.], dtype=float32)},
+     'layer2': {'weight': Array([[0., 1.]], dtype=float32),
+                'bias': Array([0.], dtype=float32)}}
+
     Args:
         tree (pytree): a pytree of arrays and scalars to ravel.
         is_leaf (callable, optional): An optionally specified function that will be called at each

optree/integration/numpy.py

@@ -46,6 +46,26 @@ def tree_ravel(
 ) -> tuple[np.ndarray, Callable[[np.ndarray], ArrayTree]]:
     r"""Ravel (flatten) a pytree of arrays down to a 1D array.
 
+    >>> tree = {
+    ...     'layer1': {
+    ...         'weight': np.arange(6, dtype=np.float32).reshape((2, 3)),
+    ...         'bias': np.arange(2, dtype=np.float32).reshape((2,)),
+    ...     },
+    ...     'layer2': {
+    ...         'weight': np.arange(2, dtype=np.float32).reshape((1, 2)),
+    ...         'bias': np.arange(1, dtype=np.float32).reshape((1,))
+    ...     },
+    ... }
+    >>> flat, unravel_func = tree_ravel(tree)
+    >>> flat
+    array([0., 1., 0., 1., 2., 3., 4., 5., 0., 0., 1.], dtype=float32)
+    >>> unravel_func(flat)
+    {'layer1': {'weight': array([[0., 1., 2.],
+                                 [3., 4., 5.]], dtype=float32),
+                'bias': array([0., 1.], dtype=float32)},
+     'layer2': {'weight': array([[0., 1.]], dtype=float32),
+                'bias': array([0.], dtype=float32)}}
+
     Args:
         tree (pytree): a pytree of arrays and scalars to ravel.
         is_leaf (callable, optional): An optionally specified function that will be called at each

optree/integration/torch.py

@@ -43,6 +43,26 @@ def tree_ravel(
 ) -> tuple[torch.Tensor, Callable[[torch.Tensor], TensorTree]]:
     r"""Ravel (flatten) a pytree of tensors down to a 1D tensor.
 
+    >>> tree = {
+    ...     'layer1': {
+    ...         'weight': torch.arange(6, dtype=torch.float64).reshape((2, 3)),
+    ...         'bias': torch.arange(2, dtype=torch.float64).reshape((2,)),
+    ...     },
+    ...     'layer2': {
+    ...         'weight': torch.arange(2, dtype=torch.float64).reshape((1, 2)),
+    ...         'bias': torch.arange(1, dtype=torch.float64).reshape((1,))
+    ...     },
+    ... }
+    >>> flat, unravel_func = tree_ravel(tree)
+    >>> flat
+    tensor([0., 1., 0., 1., 2., 3., 4., 5., 0., 0., 1.], dtype=torch.float64)
+    >>> unravel_func(flat)
+    {'layer1': {'weight': tensor([[0., 1., 2.],
+                                 [3., 4., 5.]], dtype=torch.float64),
+                'bias': tensor([0., 1.], dtype=torch.float64)},
+     'layer2': {'weight': tensor([[0., 1.]], dtype=torch.float64),
+                'bias': tensor([0.], dtype=torch.float64)}}
+
     Args:
         tree (pytree): a pytree of tensors to ravel.
         is_leaf (callable, optional): An optionally specified function that will be called at each