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