tf.experimental.dtensor.call_with_layout | TensorFlow v2.16.1 (original) (raw)

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.

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 2024-04-26 UTC.