keras/src/distribution/distribution_lib.py/DataParallel/get_variable_layout
class DataParallel:    def get_variable_layout(self, variable):
        variable_shard_spec = [None] * len(variable.shape)
        return TensorLayout(variable_shard_spec, self.device_mesh)