flax.linen.set_logical_axis_rules# flax.linen.set_logical_axis_rules(rules)[source]# Sets the global logical axis to mesh axis binding. Parameters: rules (Sequence[Tuple[str, Union[str, Tuple[str], None]]]) –