tf.experimental.dtensor.call_with_layout

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

Used in the notebooks

Used in the guide Used in the tutorials

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.

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.

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