Use MaybeJit for summaries

Author: samanklesaria

flax/nnx/rnglib.py

@@ -58,6 +58,7 @@ def _to_keyless(
 
 
 def _function_to_method(random_f):
+  @functools.wraps(random_f)
   def rngs_random_method(self: Rngs | RngStream, *args, **kwargs) -> jax.Array:
     return random_f(self(), *args, **kwargs)
 

flax/nnx/summary.py

@@ -19,6 +19,8 @@
 import io
 import typing as tp
 from types import MappingProxyType
+import functools
+from types import SimpleNamespace
 
 import jax
 import numpy as np
@@ -34,6 +36,7 @@
 
 from functools import wraps
 
+
 try:
   from IPython import get_ipython
 
@@ -50,6 +53,41 @@ class NoneDumper(yaml.SafeDumper):
   lambda dumper, data: dumper.represent_scalar('tag:yaml.org,2002:str', 'None'),
 )
 
+class MaybeJit:
+  """
+  Wraps a function with nnx.jit, but saves the original to run
+  if the function turns out to be non-concrete. We can't get the flops of non-concrete functions,
+  but we should still be able to trace the input and output shapes.
+  """
+  def __init__(self, f):
+      self.f = f
+      self.jitted = nnx.jit(f)
+      functools.update_wrapper(self, self.f)
+
+  # implement descriptor protocol so that we can use this as a method
+  def __get__(self, obj, objtype=None):
+    if obj is None:
+      return self
+    return functools.partial(self, obj)
+
+  def __call__(self, *args):
+    try:
+      return self.jitted(*args)
+    except TypeError as e:
+      return self.f(*args)
+
+  # This will only be used for the top-level method
+  def trace(self, *args):
+    return self(*args)
+
+  def lower(self, *args):
+    try:
+      return self.jitted.lower(*args)
+    except TypeError as e:
+      result = self.f(*args)
+      # Mock a `Lowered` instance with a SimpleNamespace
+      return SimpleNamespace(cost_analysis=-1, lowered=SimpleNamespace(out_info=(None, None, result)))
+
 class SizeBytes(typing.SizeBytes):
   def __repr__(self) -> str:
     bytes_repr = _bytes_repr(self.bytes)
@@ -222,7 +260,7 @@ def _overwrite_methods(env):
 
 def _get_flops(e) -> int:
   cost = e.cost_analysis() or e.compile().cost_analysis()
-  return 0 if cost is None or 'flops' not in cost else int(cost['flops'])
+  return -1 if cost is None or 'flops' not in cost else int(cost['flops'])
 
 def tabulate(
   obj,
@@ -378,7 +416,8 @@ def tabulate(
   # that each method will only be traced (and added to the table) once.
   jits = {} # Maps (class, method_name) to jit
   for key, value in saver_env.items():
-    jits[key] = nnx.jit(value)
+    jits[key] = MaybeJit(value)
+
   _overwrite_methods(jits)
 
   # Trace the top function (which indirectly traces all the others)

tests/nnx/summary_test.py

@@ -295,5 +295,17 @@ def __call__(self, x):
     self.assertEqual(module.hooked_param.get_metadata('description'), 'Custom parameter')
     self.assertEqual(module.hooked_param.get_metadata('trainable'), True)
 
+  def test_tabulate_concrete_shape(self):
+    class Net(nnx.Module):
+        def __init__(self):
+            self.rngs = nnx.Rngs(0)
+
+        def __call__(self, x):
+            return self.rngs.uniform((x.shape[0], 10))
+
+    net = Net()
+    x = jnp.zeros((4, 8))
+    print(nnx.tabulate(net, x, console_kwargs={"width": 200}))
+
 if __name__ == '__main__':
   absltest.main()