added cursor api

docs/api_reference/flax.cursor.rst

@@ -0,0 +1,60 @@
+
+flax.cursor package
+=============================
+
+The Cursor API allows for mutability of pytrees. This API provides a more
+ergonomic solution to making partial-updates of deeply nested immutable
+data structures, compared to making many nested ``dataclasses.replace`` calls.
+
+To illustrate, consider the example below::
+
+  from flax.cursor import cursor
+  import dataclasses
+  from typing import Any
+
+  @dataclasses.dataclass
+  class A:
+    x: Any
+
+  a = A(A(A(A(A(A(A(0)))))))
+
+To replace the int 0 using ``dataclasses.replace``, we would have to write many nested calls::
+
+  a2 = dataclasses.replace(
+    a,
+    x=dataclasses.replace(
+      a.x,
+      x=dataclasses.replace(
+        a.x.x,
+        x=dataclasses.replace(
+          a.x.x.x,
+          x=dataclasses.replace(
+            a.x.x.x.x,
+            x=dataclasses.replace(
+              a.x.x.x.x.x,
+              x=dataclasses.replace(a.x.x.x.x.x.x, x=1),
+            ),
+          ),
+        ),
+      ),
+    ),
+  )
+
+The equivalent can be achieved much more simply using the Cursor API::
+
+  a3 = cursor(a).x.x.x.x.x.x.x.set(1)
+  assert a2 == a3
+
+The Cursor object keeps tracks of changes made to it and when ``.build`` is called,
+generates a new object with the accumulated changes. Basic usage involves
+wrapping the object in a Cursor, making changes to the Cursor object and
+generating a new copy of the original object with the accumulated changes.
+
+.. currentmodule:: flax.cursor
+
+.. autofunction:: cursor
+
+.. autoclass:: Cursor
+  :members: build, set, apply_update
+
+

docs/api_reference/index.rst

@@ -6,6 +6,7 @@ API Reference
 
    flax.config
    flax.core.frozen_dict
+   flax.cursor
    flax.errors
    flax.jax_utils
    flax.linen/index