flax.linen.logical_axis_rules# flax.linen.logical_axis_rules(rules)[source]# Context manager for setting the logical to mesh axis bindings.
