flax.linen.override_named_call

flax.linen.override_named_call(enable=True)[source]

Returns a context manager that enables/disables named call wrapping.

Parameters

enable (bool) – If true, enables named call wrapping for labelling profile traces. (see enabled_named_call).