# Copyright 2024 The Flax Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from collections import defaultdict
import dataclasses
import inspect
import io
import typing as tp
from types import MappingProxyType
import functools
import itertools
import jax
import numpy as np
import rich.console
import rich.table
import rich.text
import yaml
import jax.numpy as jnp
from flax import nnx
from flax import typing
from flax.nnx import graphlib, statelib, variablelib
from functools import wraps
try:
from IPython import get_ipython
in_ipython = get_ipython() is not None
except ImportError:
in_ipython = False
# Custom YAML dumper to represent None as 'None' string (not YAML 'null') for clarity
class NoneDumper(yaml.SafeDumper):
pass
NoneDumper.add_representer(
type(None),
lambda dumper, data: dumper.represent_scalar('tag:yaml.org,2002:str', 'None'),
)
class SizeBytes(typing.SizeBytes):
def __repr__(self) -> str:
bytes_repr = _bytes_repr(self.bytes)
return f'{self.size:,} [dim]({bytes_repr})[/dim]'
class ObjectInfo(tp.NamedTuple):
path: statelib.PathParts
stats: dict[type[variablelib.Variable], SizeBytes]
variable_groups: defaultdict[
type[variablelib.Variable], defaultdict[typing.Key, variablelib.Variable]
]
NodeStats = dict[int, ObjectInfo | None]
def _collect_stats(
path: statelib.PathParts,
node: tp.Any,
node_stats: NodeStats,
object_types: set[type],
):
if not graphlib.is_node(node) and not isinstance(node, variablelib.Variable):
raise ValueError(f'Expected a graph node or Variable, got {type(node)!r}.')
if id(node) in node_stats:
return
stats: dict[type[variablelib.Variable], SizeBytes] = {}
variable_groups: defaultdict[
type[variablelib.Variable], defaultdict[typing.Key, variablelib.Variable]
] = defaultdict(lambda: defaultdict())
node_stats[id(node)] = ObjectInfo(path, stats, variable_groups)
if isinstance(node, nnx.Pytree):
node._nnx_tabulate_id = id(node) # type: ignore
object_types.add(type(node))
node_impl = graphlib.get_node_impl(node)
assert node_impl is not None
node_dict = node_impl.node_dict(node)
for key, value in node_dict.items():
if id(value) in node_stats:
continue
elif isinstance(value, variablelib.Variable):
var_type = type(value)
if issubclass(var_type, nnx.RngState):
var_type = nnx.RngState
size_bytes = SizeBytes.from_any(value.get_value())
if var_type in stats:
stats[var_type] += size_bytes
else:
stats[var_type] = size_bytes
variable_groups[var_type][key] = value
node_stats[id(value)] = None
elif graphlib.is_node(value):
_collect_stats((*path, key), value, node_stats, object_types)
# accumulate stats from children
child_info = node_stats[id(value)]
assert child_info is not None
for var_type, size_bytes in child_info.stats.items():
if var_type in stats:
stats[var_type] += size_bytes
else:
stats[var_type] = size_bytes
@dataclasses.dataclass(frozen=True, repr=False)
class ArrayRepr:
shape: tuple[int, ...]
dtype: tp.Any
@classmethod
def from_array(cls, x: jax.Array | np.ndarray):
return cls(jnp.shape(x), jnp.result_type(x))
def __str__(self):
shape_repr = ','.join(str(x) for x in self.shape)
return f'[dim]{self.dtype}[/dim][{shape_repr}]'
@dataclasses.dataclass
class CallInfo:
call_order: int
object_id: int
type: type
path: statelib.PathParts
inputs_repr: str
outputs: tp.Any
flops: int | None
vjp_flops: int | None
class SimpleObjectRepr:
def __init__(self, obj: tp.Any):
self.type = type(obj)
def __str__(self):
return f'{self.type.__name__}(...)'
def __repr__(self):
return f'{self.type.__name__}(...)'
def _to_dummy_array(x):
if isinstance(x,jax.ShapeDtypeStruct):
return ArrayRepr(x.shape, x.dtype)
elif isinstance(x, jax.Array | np.ndarray):
return ArrayRepr.from_array(x)
elif graphlib.is_graph_node(x):
return SimpleObjectRepr(x)
else:
return x
def _pure_nnx_vjp(f, model, *args, **kwargs):
"Wrap nnx functional api around jax.vjp. Only handles pure method calls."
graphdef, state = nnx.split(model)
def inner(state, *args, **kwargs):
model = nnx.merge(graphdef, state)
return f(model, *args, **kwargs)
return jax.vjp(inner, state, *args, **kwargs)
def filter_rng_streams(row: CallInfo):
return not issubclass(row.type, nnx.RngStream)
def _create_obj_env(object_types):
"Turn a set of object types into a dictionary mapping (type, method name) pairs to methods"
result = {}
for obj_type in object_types:
for name, top_method in inspect.getmembers(obj_type, inspect.isfunction):
if not name.startswith('_') or name == '__call__':
result[(obj_type, name)] = top_method
return result
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 _save_call_info(counter, tracer_args, f, node_stats, compute_flops, compute_vjp_flops, seen):
"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)
method_name = f.__name__
@functools.partial(jax.jit)
def jit_f(graphdef, state):
args, kwargs = nnx.merge(graphdef, state)
return f(*args, **kwargs)
@wraps(f)
def wrapper(obj, *args, **kwargs):
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)
counter_val = next(counter)
graphdef, state = nnx.split(((obj, *args), kwargs))
if compute_flops:
lowered = jit_f.lower(graphdef, state)
flops = _get_flops(lowered)
outputs = lowered.out_info
else:
flops = None
outputs = jit_f(graphdef, state)
if identifier not in seen:
seen.add(identifier)
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))
return jit_f(graphdef, state)
return wrapper
def _overwrite_methods(env):
"Overwrite methods with functions from an environment"
for (obj_type, name), f in env.items():
setattr(obj_type, name, f)
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'])
[docs]def tabulate(
obj,
*input_args,
depth: int | None = None,
method: str = '__call__',
row_filter: tp.Callable[[CallInfo], bool] = filter_rng_streams,
table_kwargs: tp.Mapping[str, tp.Any] = MappingProxyType({}),
column_kwargs: tp.Mapping[str, tp.Any] = MappingProxyType({}),
console_kwargs: tp.Mapping[str, tp.Any] = MappingProxyType({}),
compute_flops: bool = False,
compute_vjp_flops: bool = False,
**input_kwargs,
) -> str:
"""Creates a summary of the graph object represented as a table.
The table summarizes the object's state and metadata. The table is
structured as follows:
- The first column represents the path of the object in the graph.
- The second column represents the type of the object.
- The third column represents the input arguments passed to the object's
method.
- The fourth column represents the output of the object's method.
- The following columns provide information about the object's state,
grouped by Variable types.
Example::
>>> from flax import nnx
...
>>> 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))
...
>>> foo = Foo(nnx.Rngs(0))
>>> # print(nnx.tabulate(foo, jnp.ones((1, 32))))
Foo Summary
┏━━━━━━━━━━━━━━━━┳━━━━━━━━━━━┳━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━┓
┃ path ┃ type ┃ inputs ┃ outputs ┃ BatchStat ┃ Param ┃ RngState ┃
┡━━━━━━━━━━━━━━━━╇━━━━━━━━━━━╇━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━┩
│ │ Foo │ float32[1,32] │ float32[1,10] │ 276 (1.1 KB) │ 5,790 (23.2 KB) │ 2 (12 B) │
├────────────────┼───────────┼────────────────┼────────────────┼────────────────────┼─────────────────────────┼──────────┤
│ block1 │ Block │ float32[1,32] │ float32[1,128] │ 256 (1.0 KB) │ 4,480 (17.9 KB) │ 2 (12 B) │
├────────────────┼───────────┼────────────────┼────────────────┼────────────────────┼─────────────────────────┼──────────┤
│ block1/linear │ Linear │ float32[1,32] │ float32[1,128] │ │ bias: float32[128] │ │
│ │ │ │ │ │ kernel: float32[32,128] │ │
│ │ │ │ │ │ │ │
│ │ │ │ │ │ 4,224 (16.9 KB) │ │
├────────────────┼───────────┼────────────────┼────────────────┼────────────────────┼─────────────────────────┼──────────┤
│ block1/bn │ BatchNorm │ float32[1,128] │ float32[1,128] │ mean: float32[128] │ bias: float32[128] │ │
│ │ │ │ │ var: float32[128] │ scale: float32[128] │ │
│ │ │ │ │ │ │ │
│ │ │ │ │ 256 (1.0 KB) │ 256 (1.0 KB) │ │
├────────────────┼───────────┼────────────────┼────────────────┼────────────────────┼─────────────────────────┼──────────┤
│ block1/dropout │ Dropout │ float32[1,128] │ float32[1,128] │ │ │ 2 (12 B) │
├────────────────┼───────────┼────────────────┼────────────────┼────────────────────┼─────────────────────────┼──────────┤
│ block2 │ Block │ float32[1,128] │ float32[1,10] │ 20 (80 B) │ 1,310 (5.2 KB) │ │
├────────────────┼───────────┼────────────────┼────────────────┼────────────────────┼─────────────────────────┼──────────┤
│ block2/linear │ Linear │ float32[1,128] │ float32[1,10] │ │ bias: float32[10] │ │
│ │ │ │ │ │ kernel: float32[128,10] │ │
│ │ │ │ │ │ │ │
│ │ │ │ │ │ 1,290 (5.2 KB) │ │
├────────────────┼───────────┼────────────────┼────────────────┼────────────────────┼─────────────────────────┼──────────┤
│ block2/bn │ BatchNorm │ float32[1,10] │ float32[1,10] │ mean: float32[10] │ bias: float32[10] │ │
│ │ │ │ │ var: float32[10] │ scale: float32[10] │ │
│ │ │ │ │ │ │ │
│ │ │ │ │ 20 (80 B) │ 20 (80 B) │ │
├────────────────┼───────────┼────────────────┼────────────────┼────────────────────┼─────────────────────────┼──────────┤
│ block2/dropout │ Dropout │ float32[1,10] │ float32[1,10] │ │ │ │
├────────────────┼───────────┼────────────────┼────────────────┼────────────────────┼─────────────────────────┼──────────┤
│ │ │ │ Total │ 276 (1.1 KB) │ 5,790 (23.2 KB) │ 2 (12 B) │
└────────────────┴───────────┴────────────────┴────────────────┴────────────────────┴─────────────────────────┴──────────┘
Total Parameters: 6,068 (24.3 KB)
Note that ``block2/dropout`` is not shown in the table because it shares the
same ``RngState`` with ``block1/dropout``.
Args:
obj: A object to summarize. It can a pytree or a graph objects
such as nnx.Module or nnx.Optimizer.
*input_args: Positional arguments passed to the object's method.
**input_kwargs: Keyword arguments passed to the object's method.
depth: The depth of the table.
method: The method to call on the object. Default is ``'__call__'``.
row_filter: A callable that filters the rows to be displayed in the table.
By default, it filters out rows with type ``nnx.RngStream``.
table_kwargs: An optional dictionary with additional keyword arguments
that are passed to ``rich.table.Table`` constructor.
column_kwargs: An optional dictionary with additional keyword arguments
that are passed to ``rich.table.Table.add_column`` when adding columns to
the table.
console_kwargs: An optional dictionary with additional keyword arguments
that are passed to `rich.console.Console` when rendering the table.
Default arguments are ``'force_terminal': True``, and ``'force_jupyter'``
is set to ``True`` if the code is running in a Jupyter notebook, otherwise
it is set to ``False``.
compute_flops: whether to include a `flops` column in the table listing the
estimated FLOPs cost of each module forward pass. Does incur actual
on-device computation / compilation / memory allocation, but still
introduces overhead for large modules (e.g. extra 20 seconds for a
Stable Diffusion's UNet, whereas otherwise tabulation would finish in 5
seconds).
compute_vjp_flops: whether to include a `vjp_flops` column in the table
listing the estimated FLOPs cost of each module backward pass. Introduces
a compute overhead of about 2-3X of `compute_flops`.
Returns:
A string summarizing the object.
"""
_console_kwargs = {'force_terminal': True, 'force_jupyter': in_ipython}
_console_kwargs.update(console_kwargs)
obj = graphlib.clone(obj) # create copy to avoid side effects
node_stats: NodeStats = {}
object_types: set[type] = set()
_collect_stats((), obj, node_stats, object_types)
_variable_types: set[type] = {
nnx.RngState # type: ignore[misc]
if isinstance(leaf, nnx.RngState)
else type(leaf)
for _, leaf in nnx.to_flat_state(nnx.state(obj))
}
variable_types: list[type] = sorted(_variable_types, key=lambda t: t.__name__)
# Create a dictionary-version of the object's class. This makes
# iteration over methods easier.
env = _create_obj_env(object_types)
# Information is recorded in post-order, but should be presented as a pre-order traversal.
# This keeps track of the order of calls.
counter = itertools.count(0)
# Modify all the object's methods to save their lowered JIT representations.
rows : list[CallInfo] = []
seen : set = set()
jits = {k: _save_call_info(counter, rows, v, node_stats, compute_flops, compute_vjp_flops, seen)
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)
# 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)
if depth is not None:
rows = [row for row in rows if len(row.path) <= depth and row_filter(row)]
else:
rows = [row for row in rows if row_filter(row)]
rich_table = rich.table.Table(
show_header=True,
show_lines=True,
show_footer=True,
title=f'{type(obj).__name__} Summary',
**table_kwargs,
)
rich_table.add_column('path', **column_kwargs)
rich_table.add_column('type', **column_kwargs)
rich_table.add_column('inputs', **column_kwargs)
rich_table.add_column('outputs', **column_kwargs)
if compute_flops:
rich_table.add_column('flops', **column_kwargs)
if compute_vjp_flops:
rich_table.add_column('vjp_flops', **column_kwargs)
for var_type in variable_types:
rich_table.add_column(var_type.__name__, **column_kwargs)
for row in rows:
node_info = node_stats[row.object_id]
assert node_info is not None
col_reprs: list[str] = []
path_str = '/'.join(map(str, row.path))
col_reprs.append(path_str)
col_reprs.append(row.type.__name__)
col_reprs.append(row.inputs_repr)
col_reprs.append(_as_yaml_str(row.outputs))
if compute_flops:
col_reprs.append(str(row.flops))
if compute_vjp_flops:
col_reprs.append(str(row.vjp_flops))
for var_type in variable_types:
attributes = {}
variable: variablelib.Variable
for name, variable in node_info.variable_groups[var_type].items():
value = variable.get_value()
value_repr = _render_array(value) if _has_shape_dtype(value) else ''
metadata = variable.get_metadata()
for required_key in var_type.required_metadata:
metadata.pop(required_key, None)
if metadata:
attributes[name] = {
'value': value_repr,
**metadata,
}
elif value_repr:
attributes[name] = value_repr # type: ignore[assignment]
if attributes:
col_repr = _as_yaml_str(attributes) + '\n\n'
else:
col_repr = ''
size_bytes = node_info.stats.get(var_type) # type: ignore[call-overload]
if size_bytes:
col_repr += f'[bold]{size_bytes}[/bold]'
col_reprs.append(col_repr)
rich_table.add_row(*col_reprs)
total_offset = 3 + int(compute_flops) + int(compute_vjp_flops)
rich_table.columns[total_offset].footer = rich.text.Text.from_markup(
'Total', justify='right'
)
node_info = node_stats[id(obj)]
assert node_info is not None
for i, var_type in enumerate(variable_types):
size_bytes = node_info.stats[var_type]
rich_table.columns[i + total_offset + 1].footer = str(size_bytes)
rich_table.caption_style = 'bold'
total_size = sum(node_info.stats.values(), SizeBytes(0, 0))
rich_table.caption = f'\nTotal Parameters: {total_size}'
return _get_rich_repr(rich_table, _console_kwargs)
def _get_rich_repr(obj, console_kwargs):
f = io.StringIO()
console = rich.console.Console(file=f, **console_kwargs)
console.print(obj)
return f.getvalue()
def _size_and_bytes(pytree: tp.Any) -> tuple[int, int]:
leaves = jax.tree.leaves(pytree)
size = sum(x.size for x in leaves if hasattr(x, 'size'))
num_bytes = sum(
x.size * x.dtype.itemsize for x in leaves if hasattr(x, 'size')
)
return size, num_bytes
def _size_and_bytes_repr(size: int, num_bytes: int) -> str:
if not size:
return ''
bytes_repr = _bytes_repr(num_bytes)
return f'{size:,} [dim]({bytes_repr})[/dim]'
def _bytes_repr(num_bytes):
count, units = (
(f'{num_bytes / 1e9:,.1f}', 'GB')
if num_bytes > 1e9
else (f'{num_bytes / 1e6:,.1f}', 'MB')
if num_bytes > 1e6
else (f'{num_bytes / 1e3:,.1f}', 'KB')
if num_bytes > 1e3
else (f'{num_bytes:,}', 'B')
)
return f'{count} {units}'
def _has_shape_dtype(value):
return hasattr(value, 'shape') and hasattr(value, 'dtype')
def _normalize_values(x):
if isinstance(x, type):
return f'type[{x.__name__}]'
elif isinstance(x, ArrayRepr | SimpleObjectRepr):
return str(x)
else:
return repr(x)
def _maybe_pytree_to_dict(pytree: tp.Any):
path_leaves = jax.tree_util.tree_flatten_with_path(pytree)[0]
path_leaves = [
(tuple(map(graphlib._key_path_to_key, path)), value)
for path, value in path_leaves
]
if len(path_leaves) < 1:
return pytree
elif len(path_leaves) == 1 and path_leaves[0][0] == ():
return pytree
else:
return _unflatten_to_simple_structure(path_leaves, original=pytree)
def _unflatten_to_simple_structure(
xs: list[tuple[tuple[tp.Any, ...], tp.Any]], *, original: tp.Any
):
"""Rebuild a simple Python structure from path/value leaves.
This variant is aware of the original object so it can:
- Preserve empty containers that were elided by JAX flattening.
- Pad trailing missing list/tuple items using the original length.
- Distinguish placeholders for empty dict/list vs None.
"""
def _get_by_path(x, path: tuple[tp.Any, ...]):
cur = x
for k in path:
cur = cur[k]
return cur
def _to_simple(x):
# Convert to display-friendly simple structures
if isinstance(x, (list, tuple)):
return [_to_simple(e) for e in x]
if isinstance(x, dict):
return {k: _to_simple(v) for k, v in x.items()}
return x
result: list | dict = (
[] if len(xs) > 0 and isinstance(xs[0][0][0], int) else {}
)
for path, value in xs:
cursor = result
for i, key in enumerate(path[:-1]):
if isinstance(cursor, list):
# Ensure list has slot for current key; infer placeholder from original
while len(cursor) <= key:
# path to the slot we are about to create
slot_path = path[:i] + (len(cursor),)
try:
orig_slot = _get_by_path(original, slot_path)
except Exception:
orig_slot = None
if isinstance(orig_slot, (list, tuple)):
cursor.append([])
elif isinstance(orig_slot, dict):
cursor.append({})
else:
cursor.append(None)
else:
if key not in cursor:
next_key = path[i + 1]
if isinstance(next_key, int):
cursor[key] = []
else:
cursor[key] = {}
cursor = cursor[key]
if isinstance(cursor, list):
# Handle gaps in indices caused by JAX flattening eliding empty containers
while len(cursor) <= path[-1]:
slot_path = path[:-1] + (len(cursor),)
try:
orig_slot = _get_by_path(original, slot_path)
except Exception:
orig_slot = None
if isinstance(orig_slot, (list, tuple)):
cursor.append([])
elif isinstance(orig_slot, dict):
cursor.append({})
else:
cursor.append(None)
cursor[path[-1]] = value
else:
assert isinstance(cursor, dict)
cursor[path[-1]] = value
# If original is a sequence and result is a list, pad trailing items
if isinstance(original, (list, tuple)) and isinstance(result, list):
for i in range(len(result), len(original)):
slot = original[i]
result.append(_to_simple(slot))
return result
def _as_yaml_str(value) -> str:
if (hasattr(value, '__len__') and len(value) == 0) or value is None:
return ''
value = jax.tree.map(_normalize_values, value)
value = _maybe_pytree_to_dict(value)
file = io.StringIO()
yaml.dump(
value,
file,
Dumper=NoneDumper,
default_flow_style=False,
indent=2,
sort_keys=False,
explicit_end=False,
)
return file.getvalue().replace('\n...', '').replace("'", '').strip()
def _render_array(x):
shape, dtype = jnp.shape(x), jnp.result_type(x)
shape_repr = ','.join(str(x) for x in shape)
return f'[dim]{dtype}[/dim][{shape_repr}]'
def _sort_variable_types(types: tp.Iterable[type]) -> list[type]:
def _variable_parents_count(t: type):
return sum(1 for p in t.mro() if issubclass(p, variablelib.Variable))
type_sort_key = {t: (-_variable_parents_count(t), t.__name__) for t in types}
return sorted(types, key=lambda t: type_sort_key[t])
