tf.keras.distribution.distribute_tensor

Change the layout of a Tensor value in the jit function execution.

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

a new value with the specified tensor layout.