code: def call_with_layout(
      fn: Callable[..., Any],
      layout: Optional[layout_lib.Layout],
      *args,
      **kwargs,
  ) -> Any:
      """Calls a function in the DTensor device scope if `layout` is not None.
  
      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.
      """
      if layout is not None:
          if context.executing_eagerly():
              with default_mesh(layout.mesh):
                  with _dtensor_device()._default_layout(
                      layout
                  ):  # pylint: disable=protected-access
                      return fn(*args, **kwargs)
          else:
              return relayout(fn(*args, **kwargs), layout)
      return fn(*args, **kwargs)
  
imports: import contextlib
  import threading
  from typing import Any, Callable, Optional, Sequence
  from tensorflow.dtensor.python import dtensor_device
  from tensorflow.dtensor.python import gen_dtensor_ops
  from tensorflow.dtensor.python import layout
  from tensorflow.python.eager import context
  from tensorflow.python.framework import ops
  from tensorflow.python.framework import tensor
  from tensorflow.python.util import deprecation
  from tensorflow.python.util.tf_export import tf_export
label_desc: Use image node_id 1 for calling a global function with example usage: call_with_layout(fn, layout) and returns: fn, relayout, fn