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

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.

Except as otherwise noted, the content of this page is licensed under the Creative Commons Attribution 4.0 License, and code samples are licensed under the Apache 2.0 License. For details, see the Google Developers Site Policies. Java is a registered trademark of Oracle and/or its affiliates. Some content is licensed under the numpy license.

Last updated 2023-10-06 UTC.

tf.experimental.dtensor.call_with_layout

Args

Returns