tf.experimental.dtensor.call_with_layout

Calls a function in the DTensor device scope if layout is not None.

tf.experimental.dtensor.call_with_layout(
    fn: Callable[..., Any],
    layout: Optional[tf.experimental.dtensor.Layout],
    *args,
    **kwargs
) -> Any

Used in the notebooks

Used in the guide	Used in the tutorials
DTensor concepts Distributed training with Core APIs and DTensor	Distributed training with DTensors

If layout is not None, fn consumes DTensor(s) as input and produces a DTensor as output; a DTensor is a tf.Tensor with layout-related attributes.

If layout is None, fn consumes and produces regular tf.Tensors.

Args
`fn`	A supported TF API function such as tf.zeros.
`layout`	Optional, the layout of the output DTensor.
`*args`	Arguments given to `fn`.
`**kwargs`	Keyword arguments given to `fn`.

Returns
The return value of `fn` transformed to a DTensor if requested.

© 2022 The TensorFlow Authors. All rights reserved.
Licensed under the Creative Commons Attribution License 4.0.
Code samples licensed under the Apache 2.0 License.
https://www.tensorflow.org/api_docs/python/tf/experimental/dtensor/call_with_layout