flax.linen.logical_axis_rules

flax.linen.logical_axis_rules#

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

Context manager for setting the logical to mesh axis bindings.