Add tree-mode-only nnx.{vjp,jvp,jit_partial}

PiperOrigin-RevId: 874826112

Author: Cristian Garcia

docs_nnx/api_reference/flax.nnx/transforms.rst

@@ -13,6 +13,8 @@ transforms
 .. autofunction:: vmap
 .. autofunction:: eval_shape
 .. autofunction:: custom_vjp
+.. autofunction:: vjp
+.. autofunction:: jvp
 .. autofunction:: cond
 .. autofunction:: switch
 .. autofunction:: while_loop

flax/configurations.py

@@ -28,7 +28,7 @@ class Config:
   flax_max_repr_depth: int | None
   flax_always_shard_variable: bool
   flax_hijax_variable: bool
-  flax_nnx_graph_mode: bool
+  nnx_graph_mode: bool
   # See https://google.github.io/pytype/faq.html.
   _HAS_DYNAMIC_ATTRIBUTES = True
 
@@ -298,7 +298,7 @@ def static_int_env(varname: str, default: int | None) -> int | None:
   help='Whether to enable HiJAX support for `nnx.Variable`.',
 )
 nnx_graph_mode = bool_flag(
-  name='flax_nnx_graph_mode',
+  name='nnx_graph_mode',
   default=True,
   help='Whether NNX APIs default to graph-mode (True) or tree-mode (False).',
 )

flax/nnx/__init__.py

@@ -174,8 +174,11 @@
 from .transforms.autodiff import grad as grad
 from .transforms.autodiff import value_and_grad as value_and_grad
 from .transforms.autodiff import custom_vjp as custom_vjp
+from .transforms.autodiff import vjp as vjp
+from .transforms.autodiff import jvp as jvp
 from .transforms.autodiff import remat as remat
 from .transforms.compilation import jit as jit
+from .transforms.compilation import jit_partial as jit_partial
 from .transforms.compilation import shard_map as shard_map
 from .transforms.compilation import StateSharding as StateSharding
 from .transforms.iteration import Carry as Carry

flax/nnx/graph.py

@@ -769,7 +769,7 @@ def flatten(  # type: ignore[invalid-annotation]
       the overhead of the graph protocol.
   """
   if graph is None:
-    graph = config.flax_nnx_graph_mode
+    graph = config.nnx_graph_mode
   if ref_index is None:
     ref_index = RefMap()
   leaves: list[tp.Any] = []
@@ -2243,7 +2243,7 @@ def split(  # type: ignore[invalid-annotation]
     filters are passed, a single ``State`` is returned.
   """
   if graph is None:
-    graph = config.flax_nnx_graph_mode
+    graph = config.nnx_graph_mode
   graphdef, flat_state = flatten(node, graph=graph)
   flat_states = _split_state(flat_state, filters)
   states = _to_nested_state(graphdef, flat_states)
@@ -2451,7 +2451,7 @@ def state(
     One or more :class:`State` mappings.
   """
   if graph is None:
-    graph = config.flax_nnx_graph_mode
+    graph = config.nnx_graph_mode
   _, flat_state = flatten(node, graph=graph)
   state = flat_state.to_nested_state()
 
@@ -2492,7 +2492,7 @@ def graphdef(
     The :class:`GraphDef` of the :class:`Module` object.
   """
   if graph is None:
-    graph = config.flax_nnx_graph_mode
+    graph = config.nnx_graph_mode
   graphdef, _ = flatten(node, graph=graph)
   return graphdef
 
@@ -2578,7 +2578,7 @@ def pop(
     return states
 
 
-def clone(node: Node, variables: bool = True, *, graph: bool = True) -> Node:
+def clone(node: Node, variables: bool = True, *, graph: bool | None = None) -> Node:
   """Create a deep copy of the given graph node.
 
   Example usage::
@@ -2877,7 +2877,7 @@ def iter_graph(
       the overhead of the graph protocol.
   """
   if graph is None:
-    graph = config.flax_nnx_graph_mode
+    graph = config.nnx_graph_mode
   if graph:
     return _iter_graph(node)
   else:
@@ -2987,7 +2987,7 @@ def recursive_map(
       the overhead of the graph protocol.
   """
   if graph is None:
-    graph = config.flax_nnx_graph_mode
+    graph = config.nnx_graph_mode
   if graph:
     node = clone(node, variables=False)
     path_parts: PathParts = ()

flax/nnx/rnglib.py

@@ -817,7 +817,7 @@ def split_rngs(
 
   """
   if graph is None:
-    graph = config.flax_nnx_graph_mode
+    graph = config.nnx_graph_mode
 
   if isinstance(node, Missing):