remove fingerprint

PiperOrigin-RevId: 842343114

Author: Cristian Garcia

flax/nnx/graph.py

@@ -865,207 +865,6 @@ def make_mutable_arraydef(value: variablelib.Ref):
   return
 
 
-@dataclasses.dataclass(slots=True)
-class FingerprintContext:
-  next_index: int
-
-
-# TODO(cgarciae): the actual fingerprint object is not being used,
-# only the traversal process is still relevant
-def fingerprint(
-  node,
-  /,
-  *,
-  ref_index: RefMap | None = None,
-  new_ref_index: RefMap | None = None,
-) -> list[tp.Hashable]:
-  """ """
-  if ref_index is None:
-    ref_index = RefMap()
-
-  if new_ref_index is None:
-    new_ref_index = RefMap()
-  node_impl = get_node_impl(node)
-  if node_impl is None:
-    raise RuntimeError(f'Unsupported type: {type(node)}, this is a bug.')
-  ctx = FingerprintContext(len(ref_index) + len(new_ref_index))
-  fp: list[tp.Hashable] = []
-  _graph_fingerprint(ctx, fp.append, node, node_impl, ref_index, new_ref_index)
-  return fp
-
-
-def _graph_fingerprint(
-  ctx: FingerprintContext,
-  append_fn: tp.Callable[[tp.Any], None],
-  node,
-  node_impl: NodeImpl[Node, Leaf, AuxData],
-  ref_index: RefMap,
-  new_ref_index: RefMap,
-):
-  is_pytree_node_ = type(node_impl) is PytreeNodeImpl
-  is_graph_node_ = type(node_impl) is GraphNodeImpl
-
-  append_fn(type(node))
-
-  if is_graph_node_:
-    append_fn(id(node))
-    if node in ref_index:
-      append_fn(ref_index[node])
-      return
-    elif node in new_ref_index:
-      append_fn(new_ref_index[node])
-      return
-    index = new_ref_index[node] = ctx.next_index
-    ctx.next_index += 1
-  else:
-    index = -1
-
-  values, metadata = node_impl.flatten(node)
-
-  append_fn(index)
-  append_fn(metadata)
-
-  for key, value in values:
-    value_node_impl = get_node_impl(value)
-    append_fn(key)
-    if value_node_impl is not None:
-      _graph_fingerprint(
-        ctx,
-        append_fn,
-        value,
-        value_node_impl,
-        ref_index,
-        new_ref_index,
-      )
-    elif isinstance(value, Variable):
-      append_fn(id(value))
-      append_fn(type(value))
-      if value in ref_index:
-        append_fn(ref_index[value])
-      elif value in new_ref_index:
-        append_fn(new_ref_index[value])
-      else:
-        variable_index = new_ref_index[value] = ctx.next_index
-        ctx.next_index += 1
-        append_fn(variable_index)
-        for key_value in value.get_metadata().items():
-          append_fn(key_value)
-    elif not isinstance(value, (jax.Array, np.ndarray)):
-      append_fn(value)
-
-
-def check_fingerprint(
-  node,
-  fp: list[tp.Hashable],
-  /,
-  *,
-  ref_index: RefMap | None = None,
-  new_ref_index: RefMap | None = None,
-) -> bool:
-  """ """
-  if ref_index is None:
-    ref_index = RefMap()
-
-  if new_ref_index is None:
-    new_ref_index = RefMap()
-  node_impl = get_node_impl(node)
-  if node_impl is None:
-    raise RuntimeError(f'Unsupported type: {type(node)}, this is a bug.')
-  ctx = FingerprintContext(len(ref_index) + len(new_ref_index))
-  fp_matches = _check_graph_fingerprint(
-    ctx, iter(fp), node, node_impl, ref_index, new_ref_index
-  )
-  return fp_matches
-
-
-def _check_graph_fingerprint(
-  ctx: FingerprintContext,
-  fp_iterator: tp.Iterator[tp.Hashable],
-  node,
-  node_impl: NodeImpl[Node, Leaf, AuxData],
-  ref_index: RefMap,
-  new_ref_index: RefMap,
-) -> bool:
-  is_pytree_node_ = type(node_impl) is PytreeNodeImpl
-  is_graph_node_ = type(node_impl) is GraphNodeImpl
-
-  if type(node) != next(fp_iterator):
-    return False
-
-  if is_graph_node_:
-    # append_fn(id(node))
-    if id(node) != next(fp_iterator):
-      return False
-    if node in ref_index:
-      # append_fn(ref_index[node])
-      return ref_index[node] == next(fp_iterator)
-    elif node in new_ref_index:
-      # append_fn(new_ref_index[node])
-      return new_ref_index[node] == next(fp_iterator)
-    index = new_ref_index[node] = ctx.next_index
-    ctx.next_index += 1
-  else:
-    index = -1
-
-  values, metadata = node_impl.flatten(node)
-
-  # append_fn(index)
-  if index != next(fp_iterator):
-    return False
-  # append_fn(metadata)
-  if metadata != next(fp_iterator):
-    return False
-
-  for key, value in values:
-    value_node_impl = get_node_impl(value)
-    # append_fn(key)
-    if key != next(fp_iterator):
-      return False
-    if value_node_impl is not None:
-      if not _check_graph_fingerprint(
-        ctx,
-        fp_iterator,
-        value,
-        value_node_impl,
-        ref_index,
-        new_ref_index,
-      ):
-        return False
-    elif isinstance(value, Variable):
-      # append_fn(id(value))
-      if id(value) != next(fp_iterator):
-        return False
-      # append_fn(type(value))
-      if type(value) != next(fp_iterator):
-        return False
-      if value in ref_index:
-        # append_fn(ref_index[value])
-        if ref_index[value] != next(fp_iterator):
-          return False
-      elif value in new_ref_index:
-        # append_fn(new_ref_index[value])
-        if new_ref_index[value] != next(fp_iterator):
-          return False
-      else:
-        variable_index = new_ref_index[value] = ctx.next_index
-        ctx.next_index += 1
-        # append_fn(variable_index)
-        if variable_index != next(fp_iterator):
-          return False
-        for key_value in value.get_metadata().items():
-          # append_fn(key_value)
-          if key_value != next(fp_iterator):
-            return False
-    else:
-      if isinstance(value, (jax.Array, np.ndarray)):
-        raise ValueError(f'Arrays leaves are not supported: {value}')
-      # append_fn(value)
-      if value != next(fp_iterator):
-        return False
-
-  return True
-
-
 def _get_sorted_leaves(
   xs: tp.Mapping[tp.Any, tp.Any],
 ) -> list[tp.Any]:
@@ -1629,13 +1428,17 @@ def create_static_cache(x):
       node_cache = unflatten(
         graphdef, flat_state, index_ref=index_ref, copy_variables=False
       )
-      cached_new_ref_index = RefMap()
-      _fp = fingerprint(
+      start_index = len(cached_ref_index)
+      flatten(
         node_cache,
         ref_index=cached_ref_index,
-        new_ref_index=cached_new_ref_index,
+        with_paths=False,
+      )
+      cached_new_ref_index = RefMap(
+        (key, value)
+        for key, value in cached_ref_index.items()
+        if value >= start_index
       )
-      cached_ref_index.update(cached_new_ref_index)
       cache[node_cache] = StaticCache.create(
         graphdef, paths, variables, cached_new_ref_index
       )

tests/nnx/graph_utils_test.py

@@ -845,44 +845,6 @@ def f(*pure_args):
     self.assertIs(m1, args_out[2]['b'])
     self.assertIs(m2, args_out[1])
 
-  def test_fingerprint_basic(self):
-    m = nnx.Linear(2, 3, rngs=nnx.Rngs(0))
-    fp1 = nnx.graph.fingerprint(m)
-    fp2 = nnx.graph.fingerprint(m)
-
-    self.assertEqual(fp1, fp2)
-    self.assertTrue(nnx.graph.check_fingerprint(m, fp1))
-    self.assertTrue(nnx.graph.check_fingerprint(m, fp2))
-
-  def test_fingerprint_variable_id_sensitive(self):
-    m1 = nnx.Linear(2, 3, rngs=nnx.Rngs(0))
-    fp1 = nnx.graph.fingerprint(m1)
-
-    m2 = nnx.Linear(2, 3, rngs=nnx.Rngs(0))
-    fp2 = nnx.graph.fingerprint(m2)
-
-    self.assertNotEqual(fp1, fp2)
-    self.assertTrue(nnx.graph.check_fingerprint(m1, fp1))
-    self.assertTrue(nnx.graph.check_fingerprint(m2, fp2))
-    self.assertFalse(nnx.graph.check_fingerprint(m1, fp2))
-    self.assertFalse(nnx.graph.check_fingerprint(m2, fp1))
-
-  def test_fingerprint_module_id_insensitive(self):
-    m1 = nnx.Linear(2, 3, rngs=nnx.Rngs(0))
-    m2 = nnx.Linear(2, 3, rngs=nnx.Rngs(0))
-
-    m1.kernel = m2.kernel
-    m1.bias = m2.bias
-
-    fp1 = nnx.graph.fingerprint(m1)
-    fp2 = nnx.graph.fingerprint(m2)
-
-    self.assertNotEqual(fp1, fp2)
-    self.assertTrue(nnx.graph.check_fingerprint(m1, fp1))
-    self.assertTrue(nnx.graph.check_fingerprint(m2, fp2))
-    self.assertFalse(nnx.graph.check_fingerprint(m1, fp2))
-    self.assertFalse(nnx.graph.check_fingerprint(m2, fp1))
-
   def test_split_variable(self):
     v = nnx.Param(1)
     graphdef, state = nnx.split(v)
@@ -1137,6 +1099,27 @@ def test_iter_graph(self):
     actual = [node for node in nodes if any(node is e for e in unique)]
     self.assertEqual(list(map(index, actual)), list(map(index, expected)))
 
+  def test_cached_partial_docstring_example(self):
+    import optax
+
+    model = nnx.Linear(2, 3, rngs=nnx.Rngs(0))
+    optimizer = nnx.Optimizer(model, optax.adamw(1e-3), wrt=nnx.Param)
+
+    @nnx.jit
+    def train_step(model, optimizer, x, y):
+      def loss_fn(model):
+        return jnp.mean((model(x) - y) ** 2)
+
+      loss, grads = nnx.value_and_grad(loss_fn)(model)
+      optimizer.update(model, grads)
+      return loss
+
+    cached_train_step = nnx.cached_partial(train_step, model, optimizer)
+
+    for step in range(2):
+      x, y = jnp.ones((10, 2)), jnp.ones((10, 3))
+      loss = cached_train_step(x, y)
+      self.assertIsInstance(loss, jax.Array)
 
   def test_find_duplicates(self):
     class SharedModules(nnx.Module):