keras/src/backend/jax/trainer.py/_distribute_data
def _distribute_data(data):
    distribution = distribution_lib.distribution()
    if distribution is not None:

        def distribute_single_value(d):
            layout = distribution.get_data_layout(d.shape)
            return jax_distribution_lib.distribute_data_input(d, layout)

        return tree.map_structure(distribute_single_value, data)
    else:
        return tree.map_structure(jax.device_put, data)