flax.linen.compact

flax.linen.compact(fun)

Marks the given module method allowing inlined submodules.

Methods wrapped in @compact can define submodules directly within the method.

For instance:

@compact
__call__(self, x, features):
  x = nn.Dense(features)(x)
  ...

At most one method in each Module may be wrapped with @compact.

Parameters

fun – The Module method to mark as compact.

Returns

The given function fun marked as compact.