Skip to main content
Ctrl+K

flax.linen.logical_axis_rules#

flax.linen.logical_axis_rules(rules)[source]#

Context manager for setting the logical to mesh axis bindings.

Parameters:

rules (Sequence[Tuple[str, Union[str, Tuple[str], None]]]) –