def distribute_tensor(tensor, layout):
    """Change the layout of a Tensor value in the jit function execution.

    Args:
        tensor: a Tensor to change the layout.
        layout: `TensorLayout` to be applied on the value.

    Returns:
        a new value with the specified tensor layout.
    """
    if isinstance(tensor, KerasTensor):
        # keras tensor is only used for building functional model, and can't be
        # used to alter layout/sharding.
        return tensor
    return distribution_lib.distribute_tensor(tensor, layout)
Change the layout of a Tensor value in the jit function execution.

    Args:
        tensor: a Tensor to change the layout.
        layout: `TensorLayout` to be applied on the value.

    Returns:
        a new value with the specified tensor layout.
    
keras/src/distribution/distribution_lib.py