Clean up code

Author: samanklesaria

flax/nnx/summary.py

@@ -20,6 +20,7 @@
 import typing as tp
 from types import MappingProxyType
 import functools
+import itertools
 from types import SimpleNamespace
 
 import jax
@@ -83,7 +84,7 @@ def lower(self, *args, **kwargs):
     except TypeError as e:
       result = self.f(*args, **kwargs)
       # Mock a `Lowered` instance with a SimpleNamespace
-      return SimpleNamespace(cost_analysis=-1, lowered=SimpleNamespace(out_info=(None, None, result)))
+      return SimpleNamespace(cost_analysis="0", lowered=SimpleNamespace(out_info=(None, None, result)))
 
 class SizeBytes(typing.SizeBytes):
   def __repr__(self) -> str:
@@ -165,6 +166,7 @@ def __str__(self):
 
 @dataclasses.dataclass
 class CallInfo:
+  call_order: int
   object_id: int
   type: type
   path: statelib.PathParts
@@ -201,28 +203,6 @@ def inner(state, *args, **kwargs):
     return f(model, *args, **kwargs)
   return jax.vjp(inner, state, *args, **kwargs)
 
-def _get_call_info(lowered, method_name, node_stats, obj, compute_flops, inputs_repr, vjp_flops):
-  flops = _get_flops(lowered) if compute_flops else None
-  outputs = lowered.lowered.out_info[2]
-  output_repr = jax.tree.map(_to_dummy_array, outputs)
-  object_id: int = getattr(obj, '_nnx_tabulate_id')
-  node_info = node_stats[object_id]
-  assert node_info is not None
-  path = node_info.path
-  if method_name != '__call__':
-    path = (*path, method_name)
-
-  return CallInfo(
-    object_id=object_id,
-    type=type(obj),
-    path=path,
-    inputs_repr=inputs_repr,
-    outputs=output_repr,
-    flops=flops,
-    vjp_flops=vjp_flops
-  )
-
-
 def filter_rng_streams(row: CallInfo):
   return not issubclass(row.type, nnx.RngStream)
 
@@ -235,38 +215,52 @@ def _create_obj_env(object_types):
         result[(obj_type, name)] = top_method
   return result
 
-def _argsave(counter, tracer_args, f, compute_vjp_flops):
+def _get_inputs_repr(args, kwargs):
+  input_args, input_kwargs = jax.tree.map(
+    _to_dummy_array, (args, kwargs)
+  )
+  inputs_repr = ''
+  if input_args:
+    if len(input_args) == 1 and not input_kwargs:
+      inputs_repr += _as_yaml_str(input_args[0])
+    else:
+      inputs_repr += _as_yaml_str(input_args)
+    if input_kwargs:
+      inputs_repr += '\n'
+  if input_kwargs:
+    inputs_repr += _as_yaml_str(input_kwargs)
+  return inputs_repr
+
+def _argsave(counter, tracer_args, f, node_stats, compute_flops, compute_vjp_flops):
   "Wrap a function to save its arguments"
+
+  # Used when computing vjp flops
   def do_vjp(*args, **kwargs):
     primals, f_vjp = jax.vjp(f, *args, **kwargs)
     return f_vjp(primals)
-  n = f.__name__
+
+  method_name = f.__name__
+
   @wraps(f)
   def wrapper(obj, *args, **kwargs):
-    input_args, input_kwargs = jax.tree.map(
-      _to_dummy_array, (args, kwargs)
-    )
-    inputs_repr = ''
-    if input_args:
-      if len(input_args) == 1 and not input_kwargs:
-        inputs_repr += _as_yaml_str(input_args[0])
-      else:
-        inputs_repr += _as_yaml_str(input_args)
-      if input_kwargs:
-        inputs_repr += '\n'
-    if input_kwargs:
-      inputs_repr += _as_yaml_str(input_kwargs)
-
-    identifier = (inputs_repr, getattr(obj, '_nnx_tabulate_id'))
+    inputs_repr = _get_inputs_repr(args, kwargs)
+    object_id = getattr(obj, '_nnx_tabulate_id')
+    node_info = node_stats[object_id]
+    path = node_info.path
+    if method_name != '__call__':
+      path = (*path, method_name)
+    identifier = (inputs_repr, object_id)
     if identifier not in f.seen:
-      counter_val = counter[0]
-      counter[0] += 1
+      counter_val = next(counter)
       lowered = f.lower(obj, *args, **kwargs)
-      if compute_vjp_flops:
-        vjp_flops = _get_flops(jax.jit(do_vjp).lower(obj, *args, **kwargs))
-      else:
-        vjp_flops = None
-      tracer_args.append((counter_val, obj, n, lowered, inputs_repr, vjp_flops))
+      flops = _get_flops(lowered) if compute_flops else None
+      outputs = lowered.lowered.out_info[2]
+      output_repr = jax.tree.map(_to_dummy_array, outputs)
+      vjp_flops = _get_flops(jax.jit(do_vjp).lower(
+        obj, *args, **kwargs)) if compute_vjp_flops else None
+      tracer_args.append(
+        CallInfo(counter_val, object_id, type(obj), path, inputs_repr,
+          output_repr, flops, vjp_flops))
       f.seen.add(identifier)
     return f(obj, *args, **kwargs)
   return wrapper
@@ -278,7 +272,7 @@ def _overwrite_methods(env):
 
 def _get_flops(e) -> int:
   cost = e.cost_analysis() or e.compile().cost_analysis()
-  return -1 if cost is None or 'flops' not in cost else int(cost['flops'])
+  return 0 if cost is None or 'flops' not in cost else int(cost['flops'])
 
 def tabulate(
   obj,
@@ -424,22 +418,20 @@ def tabulate(
   env = _create_obj_env(object_types)
 
   # Information is recorded in post-order, but should be presented as a pre-order traversal.
-  # This counter is incremented in pre-order traversal to keep track of the order of calls.
-  counter = [0]
+  # This keeps track of the order of calls.
+  counter = itertools.count(0)
 
   # Modify all the object's methods to save their lowered JIT representations.
-  tracer_args = []
-  jits = {k: _argsave(counter, tracer_args, MaybeJit(v), compute_vjp_flops) for k,v in env.items()}
+  rows : list[CallInfo] = []
+  jits = {k: _argsave(counter, rows, MaybeJit(v), node_stats, compute_flops, compute_vjp_flops)
+    for k,v in env.items()}
   _overwrite_methods(jits)
 
   # Trace the top function (which indirectly traces all the others)
   jits[(type(obj), method)](obj, *input_args, **input_kwargs)
 
-  # Get call_info
-  rows : list[CallInfo] = [_get_call_info(
-    lowered, name, node_stats, object,
-    compute_flops, inputs_repr, vjp_flops)
-    for (_, object, name, lowered, inputs_repr, vjp_flops) in sorted(tracer_args, key=lambda x: x[0])]
+  # Sort call info in pre-order traversal order
+  rows.sort(key=lambda x: x.call_order)
 
   # Restore the object's original methods
   _overwrite_methods(env)

tests/nnx/summary_test.py

@@ -308,7 +308,5 @@ def __call__(self, x):
     x = jnp.zeros((4, 8))
     nnx.tabulate(net, x, console_kwargs={"width": 200})
 
-# TODO: should test dynamic shapes with nested calls. This will probably lead to duplicate rows.
-
 if __name__ == '__main__':
   absltest.main()