keras/src/distribution/distribution_lib.py/TensorLayout/__init__
class TensorLayout:    def __init__(self, axes, device_mesh=None):
        self._axes = tuple(axes)
        self._device_mesh = device_mesh
        self._validate_axes()