Docs4dev
Docs
tensorflow / 2.9 / experimental / dtensor / call_with_layout.html

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.

© 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/versions/r2.9/api_docs/python/tf/experimental/dtensor/call_with_layout