flax.linen.override_named_call

flax.linen.override_named_call(enable=True)

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

Parameters

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