flax.linen.set_logical_axis_rules# flax.linen.set_logical_axis_rules(rules)[source]# Sets the global logical axis to mesh axis binding.