View source on GitHub |
A layout to apply to a tensor.
tf.keras.distribution.TensorLayout(
axes, device_mesh=None
)
This API is aligned with jax.sharding.NamedSharding
and tf.dtensor.Layout
.
See more details in jax.sharding.NamedSharding and tf.dtensor.Layout.
Attributes | |
---|---|
axes
|
|
device_mesh
|