flax.linen.set_logical_axis_rules#

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

Sets the global logical axis to mesh axis binding.