Adding set_mode_info

Author: chapman20j

flax/nnx/__init__.py

@@ -48,6 +48,7 @@
 from .module import M as M
 from .module import Module as Module
 from .module import set_mode as set_mode
+from .module import set_mode_info as set_mode_info
 from .module import train_mode as train_mode
 from .module import eval_mode as eval_mode
 from .module import iter_children as iter_children, iter_modules as iter_modules

flax/nnx/module.py

@@ -14,6 +14,7 @@
 
 from __future__ import annotations
 
+import inspect
 import typing as tp
 
 import jax
@@ -559,6 +560,115 @@ def eval_mode(node: A, /, *, only: filterlib.Filter = ..., **kwargs) -> A:
   )
 
 
+def _parse_docstring_args(doc_str: str) -> dict[str, str]:
+  """Parses parameters from `Args:` section of a function docstring.
+  Assumes Google style docstrings. Returns a dictionary with
+  keys representing argument names and values representing descriptions.
+  Each description has lines starting with 4 spaces.
+  """
+  lines = doc_str.split("\n")
+
+  # Get lines with the parameter names
+  inds = [i for i, l in enumerate(lines) if l.startswith("  ") and not l.startswith("    ")]
+  inds.append(len(lines))
+  out = dict()
+
+  # Parse each argument
+  for i in range(len(inds)-1):
+    start, end = inds[i], inds[i+1]
+
+    # Process first line for the description
+    first_colon = lines[start].find(":")
+    name = lines[start][:first_colon].strip()
+    desc = [" "*4 + lines[start][first_colon+1:].strip()]
+
+    # Append remaining description lines
+    for j in range(start+1, end):
+      desc.append(lines[j])
+
+    out[name] = "\n".join(desc)
+  return out
+
+
+
+def set_mode_info(node: Module, /, *, only: filterlib.Filter = ...) -> str:
+  """Provides information about the ``set_mode`` arguments for a module and all
+  submodules. If no docstring is provided for a module's `set_mode`, this function
+  puts the `set_mode` signature below the function.
+
+  Example::
+    >>> from flax import nnx
+    ...
+    >>> class CustomModel(nnx.Module):
+    ...   def __init__(self, *, rngs):
+    ...       self.mha = nnx.MultiHeadAttention(4, 8, 32, rngs=rngs)
+    ...       self.drop = nnx.Dropout(0.5, rngs=rngs)
+    ...       self.bn = nnx.BatchNorm(32, rngs=rngs)
+    ...
+    >>> model = CustomModel(rngs=nnx.Rngs(0))
+    >>> print(nnx.set_mode_info(model))
+    BatchNorm:
+      use_running_average: bool | None = None
+        if True, the stored batch statistics will be
+        used instead of computing the batch statistics on the input.
+    Dropout:
+      deterministic: bool | None = None
+        if True, disables dropout masking.
+    MultiHeadAttention:
+      deterministic: bool | None = None
+        if True, the module is set to deterministic mode.
+      decode: bool | None = None
+        if True, the module is set to decode mode.
+      batch_size: int | Shape | None = None
+        the batch size to use for the cache.
+      max_length: int | None = None
+        the max length to use for the cache.
+
+  Args:
+    node: the object to display ``set_mode`` information for.
+    only: Filters to select the Modules to display information for.
+  """
+  predicate = filterlib.to_predicate(only)
+  classes: set[Module] = set()
+
+  def _set_mode_info_fn(path, node):
+    if hasattr(node, 'set_mode') and predicate(path, node):
+      classes.add(node.__class__)
+    return node
+
+  graph.recursive_map(_set_mode_info_fn, node)
+
+  class_list = sorted(list(classes), key=lambda x: x.__qualname__)
+  out_str = []
+  for c in class_list:
+    out_str.append(f"{c.__qualname__}:")
+    sig = inspect.signature(c.set_mode)
+    doc = inspect.getdoc(c.set_mode)
+
+    # Parse docstring
+    if isinstance(doc, str):
+      start, end = doc.find("Args:\n"), doc.find("Returns:\n")
+      if end == -1:
+        end = len(doc)
+      doc = doc[start+6:end]
+      parsed_docstring = _parse_docstring_args(doc)
+
+      # Generate output from signature and docstring
+      skip_names = {"self", "args", "kwargs"}
+      for name, param in sig.parameters.items():
+        if name in skip_names:
+          continue
+
+        if param.default is inspect.Parameter.empty:
+          out_str.append(f"  {name}: {param.annotation}")
+        else:
+          out_str.append(f"  {name}: {param.annotation} = {param.default}")
+        out_str.append(parsed_docstring[name])
+    else:
+      out_str.append(f"  set_mode{sig}")
+
+
+  return "\n".join(out_str)
 
 def first_from(*args: tp.Optional[A], error_msg: str) -> A:
   """Return the first non-None argument.

tests/nnx/module_test.py

@@ -759,6 +759,69 @@ def __call__(self, x):
     self.assertIn(str(expected_total_batch_stats), foo_repr[0])
     self.assertIn(str(expected_total_rng_states), foo_repr[0])
 
+  def test_set_mode_info(self):
+    class Block(nnx.Module):
+      def __init__(self, din, dout, *, rngs: nnx.Rngs):
+        self.linear = nnx.Linear(din, dout, rngs=rngs)
+        self.bn = nnx.BatchNorm(dout, rngs=rngs)
+        self.dropout = nnx.Dropout(0.2, rngs=rngs)
+
+      def __call__(self, x):
+        return nnx.relu(self.dropout(self.bn(self.linear(x))))
+
+    class Foo(nnx.Module):
+      def __init__(self, rngs: nnx.Rngs):
+        self.block1 = Block(32, 128, rngs=rngs)
+        self.block2 = Block(128, 10, rngs=rngs)
+
+      def __call__(self, x):
+        return self.block2(self.block1(x))
+
+    obj = Foo(rngs=nnx.Rngs(0))
+    info_str = nnx.set_mode_info(obj)
+    self.assertEqual(info_str.count("BatchNorm:"), 1)
+    self.assertEqual(info_str.count("Dropout:"), 1)
+
+  def test_set_mode_info_with_filter(self):
+    class Block(nnx.Module):
+      def __init__(self, din, dout, *, rngs: nnx.Rngs):
+        self.linear = nnx.Linear(din, dout, rngs=rngs)
+        self.bn = nnx.BatchNorm(dout, rngs=rngs)
+        self.dropout = nnx.Dropout(0.2, rngs=rngs)
+
+      def __call__(self, x):
+        return nnx.relu(self.dropout(self.bn(self.linear(x))))
+
+    obj = Block(4, 8, rngs=nnx.Rngs(0))
+    info_str = nnx.set_mode_info(obj, only=nnx.Dropout)
+    self.assertIn("Dropout:", info_str)
+    self.assertNotIn("BatchNorm:", info_str)
+
+    info_str = nnx.set_mode_info(obj, only=nnx.MultiHeadAttention)
+    self.assertEmpty(info_str)
+
+  def test_set_mode_info_with_custom_set_mode(self):
+    class Block(nnx.Module):
+      def __init__(self, *, rngs: nnx.Rngs):
+        pass
+
+      def __call__(self, x):
+        return x
+
+      def set_mode(self, arg1: bool | None = None, arg2: int | None = None, **kwargs) -> dict:
+        """Example set mode test. This follows Google style docstrings.
+
+        Args:
+          arg1: The first argument.
+          arg2: The second argument.
+            This has two lines.
+        """
+        return kwargs
+
+    obj = Block(rngs=nnx.Rngs(0))
+    info_str = nnx.set_mode_info(obj)
+    self.assertEqual(f"{obj.__class__.__qualname__}:\n  arg1: bool | None = None\n    The first argument.\n  arg2: int | None = None\n    The second argument.\n    This has two lines.", info_str)
+
 
 class TestModuleDataclass:
   def test_basic(self):