ott.solvers.nn.models.MLP.tabulate#

MLP.tabulate(rngs, *args, depth=None, show_repeated=False, mutable=True, console_kwargs=None, **kwargs)#

Creates a summary of the Module represented as a table.

This method has the same signature and internally calls Module.init, but instead of returning the variables, it returns the string summarizing the Module in a table. tabulate uses jax.eval_shape to run the forward computation without consuming any FLOPs or allocating memory.

Additional arguments can be passed into the console_kwargs argument, for example, {‘width’: 120}. For a full list of console_kwargs arguments, see: https://rich.readthedocs.io/en/stable/reference/console.html#rich.console.Console

Example:

import jax
import jax.numpy as jnp
import flax.linen as nn

class Foo(nn.Module):
    @nn.compact
    def __call__(self, x):
        h = nn.Dense(4)(x)
        return nn.Dense(2)(h)

x = jnp.ones((16, 9))

print(Foo().tabulate(jax.random.PRNGKey(0), x))

This gives the following output:

                                Foo Summary
┏━━━━━━━━━┳━━━━━━━━┳━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━┓
┃ path    ┃ module ┃ inputs        ┃ outputs       ┃ params               ┃
┡━━━━━━━━━╇━━━━━━━━╇━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━┩
│         │ Foo    │ float32[16,9] │ float32[16,2] │                      │
├─────────┼────────┼───────────────┼───────────────┼──────────────────────┤
│ Dense_0 │ Dense  │ float32[16,9] │ float32[16,4] │ bias: float32[4]     │
│         │        │               │               │ kernel: float32[9,4] │
│         │        │               │               │                      │
│         │        │               │               │ 40 (160 B)           │
├─────────┼────────┼───────────────┼───────────────┼──────────────────────┤
│ Dense_1 │ Dense  │ float32[16,4] │ float32[16,2] │ bias: float32[2]     │
│         │        │               │               │ kernel: float32[4,2] │
│         │        │               │               │                      │
│         │        │               │               │ 10 (40 B)            │
├─────────┼────────┼───────────────┼───────────────┼──────────────────────┤
│         │        │               │         Total │ 50 (200 B)           │
└─────────┴────────┴───────────────┴───────────────┴──────────────────────┘

                      Total Parameters: 50 (200 B)

Note: rows order in the table does not represent execution order, instead it aligns with the order of keys in variables which are sorted alphabetically.

Parameters:

rngs (Union[Any, Dict[str, Any]]) – The rngs for the variable collections as passed to Module.init.
*args – The arguments to the forward computation.
depth (Optional[int]) – controls how many submodule deep the summary can go. By default its None which means no limit. If a submodule is not shown because of the depth limit, its parameter count and bytes will be added to the row of its first shown ancestor such that the sum of all rows always adds up to the total number of parameters of the Module.
show_repeated (bool) – If True, repeated calls to the same module will be shown in the table, otherwise only the first call will be shown. Default is False.
mutable (Union[bool, str, Collection[str], DenyList]) – Can be bool, str, or list. Specifies which collections should be treated as mutable: bool: all/no collections are mutable. str: The name of a single mutable collection. list: A list of names of mutable collections. By default all collections except ‘intermediates’ are mutable.
console_kwargs (Optional[Mapping[str, Any]]) – An optional dictionary with additional keyword arguments that are passed to rich.console.Console when rendering the table. Default arguments are {‘force_terminal’: True, ‘force_jupyter’: False}.
**kwargs – keyword arguments to pass to the forward computation.

Return type:

str

Returns:

A string summarizing the Module.

ott.solvers.nn.models.MLP.tabulate

Contents

ott.solvers.nn.models.MLP.tabulate#