Merge pull request #4632 from google:add-max-repr-depth

Flax Authors · Flax Authors · commit dd3c18be0f40 · 2025-03-17T18:04:54.000-07:00
PiperOrigin-RevId: 737815908
diff --git a/flax/configurations.py b/flax/configurations.py
@@ -23,6 +23,7 @@
 
 class Config:
   flax_use_flaxlib: bool
+  flax_max_repr_depth: int | None
   # See https://google.github.io/pytype/faq.html.
   _HAS_DYNAMIC_ATTRIBUTES = True
 
@@ -122,6 +123,38 @@ def bool_flag(name: str, *, default: bool, help: str) -> FlagHolder[bool]:
   return fh
 
 
+def int_flag(name: str, *, default: int | None, help: str) -> FlagHolder[int]:
+  """Set up an integer flag.
+
+  Example::
+
+    num_foo = int_flag(
+        name='flax_num_foo',
+        default=42,
+        help='Number of foo.',
+    )
+
+  Now the ``FLAX_NUM_FOO`` shell environment variable can be used to
+  control the process-level value of the flag, in addition to using e.g.
+  ``config.update("flax_num_foo", 42)`` directly.
+
+  Args:
+    name: converted to lowercase to define the name of the flag. It is
+      converted to uppercase to define the corresponding shell environment
+      variable.
+    default: a default value for the flag.
+    help: used to populate the docstring of the returned flag holder object.
+
+  Returns:
+    A flag holder object for accessing the value of the flag.
+  """
+  name = name.lower()
+  config._add_option(name, static_int_env(name.upper(), default))
+  fh = FlagHolder[int](name, help)
+  setattr(Config, name, property(lambda _: fh.value, doc=help))
+  return fh
+
+
 def static_bool_env(varname: str, default: bool) -> bool:
   """Read an environment variable and interpret it as a boolean.
 
@@ -149,6 +182,27 @@ def static_bool_env(varname: str, default: bool) -> bool:
     )
 
 
+def static_int_env(varname: str, default: int | None) -> int | None:
+  """Read an environment variable and interpret it as an integer.
+
+  Args:
+    varname: the name of the variable
+    default: the default integer value
+  Returns:
+    integer return value derived from defaults and environment.
+  Raises: ValueError if the environment variable is not an integer.
+  """
+  val = os.getenv(varname)
+  if val is None:
+    return default
+  try:
+    return int(val)
+  except ValueError:
+    raise ValueError(
+      f'invalid integer value {val!r} for environment {varname!r}'
+    ) from None
+
+
 @contextmanager
 def temp_flip_flag(var_name: str, var_value: bool):
   """Context manager to temporarily flip feature flags for test functions.
@@ -211,4 +265,10 @@ def temp_flip_flag(var_name: str, var_value: bool):
   name='flax_use_flaxlib',
   default=False,
   help='Whether to use flaxlib for C++ acceleration.',
+)
+
+flax_max_repr_depth = int_flag(
+  name='flax_max_repr_depth',
+  default=None,
+  help='Maximum depth of reprs for nested flax objects. Default is None (no limit).',
 )
diff --git a/flax/nnx/reprlib.py b/flax/nnx/reprlib.py
@@ -17,6 +17,7 @@
 import sys
 import threading
 import typing as tp
+from flax import config as flax_config
 
 A = tp.TypeVar('A')
 B = tp.TypeVar('B')
@@ -90,6 +91,7 @@ class Color(tp.NamedTuple):
 @dataclasses.dataclass
 class ReprContext(threading.local):
   current_color: Color = COLOR
+  depth: int = 0
 
 
 REPR_CONTEXT = ReprContext()
@@ -172,51 +174,62 @@ def __str__(self) -> str:
 
 
 def get_repr(obj: Representable) -> str:
-  if not isinstance(obj, Representable):
-    raise TypeError(f'Object {obj!r} is not representable')
-
-  c = REPR_CONTEXT.current_color
-  iterator = obj.__nnx_repr__()
-  config = next(iterator)
+  REPR_CONTEXT.depth += 1
+  try:
+    if not isinstance(obj, Representable):
+      raise TypeError(f'Object {obj!r} is not representable')
 
-  if not isinstance(config, Object):
-    raise TypeError(f'First item must be Config, got {type(config).__name__}')
+    c = REPR_CONTEXT.current_color
+    iterator = obj.__nnx_repr__()
+    config = next(iterator)
 
-  kv_sep = f'{c.SEP}{config.kv_sep}{c.END}'
+    if not isinstance(config, Object):
+      raise TypeError(f'First item must be Config, got {type(config).__name__}')
 
-  def _repr_elem(elem: tp.Any) -> str:
-    if not isinstance(elem, Attr):
-      raise TypeError(f'Item must be Elem, got {type(elem).__name__}')
+    kv_sep = f'{c.SEP}{config.kv_sep}{c.END}'
 
-    value_repr = elem.value if elem.use_raw_value else colorized(elem.value)
-    value_repr = value_repr.replace('\n', '\n' + config.indent)
-    key = elem.key if elem.use_raw_key else f'{c.ATTRIBUTE}{elem.key}{c.END}'
-    indent = '' if config.same_line else config.indent
+    def _repr_elem(elem: tp.Any) -> str:
+      if not isinstance(elem, Attr):
+        raise TypeError(f'Item must be Elem, got {type(elem).__name__}')
 
-    return f'{indent}{elem.start}{key}{kv_sep}{value_repr}{elem.end}'
+      value_repr = elem.value if elem.use_raw_value else colorized(elem.value)
+      value_repr = value_repr.replace('\n', '\n' + config.indent)
+      key = elem.key if elem.use_raw_key else f'{c.ATTRIBUTE}{elem.key}{c.END}'
+      indent = '' if config.same_line else config.indent
 
-  elems = config.elem_sep.join(map(_repr_elem, iterator))
+      return f'{indent}{elem.start}{key}{kv_sep}{value_repr}{elem.end}'
 
-  if elems:
-    if config.same_line:
-      elems_repr = elems
-      comment = ''
+    max_depth_reached = (
+      flax_config.flax_max_repr_depth is not None
+      and REPR_CONTEXT.depth > flax_config.flax_max_repr_depth
+    )
+    if max_depth_reached:
+      elems = '...'
     else:
-      elems_repr = '\n' + elems + '\n'
-      comment = f'{c.COMMENT}{config.comment}{c.END}'
-  else:
-    elems_repr = config.empty_repr
-    comment = ''
-
-  type_repr = (
-    config.type if isinstance(config.type, str) else config.type.__name__
-  )
-  type_repr = f'{c.TYPE}{type_repr}{c.END}' if type_repr else ''
-  start = f'{c.PAREN}{config.start}{c.END}' if config.start else ''
-  end = f'{c.PAREN}{config.end}{c.END}' if config.end else ''
+      elems = config.elem_sep.join(map(_repr_elem, iterator))
+
+    if elems:
+      if config.same_line or max_depth_reached:
+        elems_repr = elems
+        comment = ''
+      else:
+        elems_repr = '\n' + elems + '\n'
+        comment = f'{c.COMMENT}{config.comment}{c.END}'
+    else:
+      elems_repr = config.empty_repr
+      comment = ''
 
-  out = f'{type_repr}{start}{comment}{elems_repr}{end}'
-  return out
+    type_repr = (
+      config.type if isinstance(config.type, str) else config.type.__name__
+    )
+    type_repr = f'{c.TYPE}{type_repr}{c.END}' if type_repr else ''
+    start = f'{c.PAREN}{config.start}{c.END}' if config.start else ''
+    end = f'{c.PAREN}{config.end}{c.END}' if config.end else ''
+
+    out = f'{type_repr}{start}{comment}{elems_repr}{end}'
+    return out
+  finally:
+    REPR_CONTEXT.depth -= 1
 
 class MappingReprMixin(Representable):
   def __nnx_repr__(self):
