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.