code: def np_function(func=None, output_dtypes=None):
      """Decorator that allow a numpy function to be used in Eager and Graph modes.
  
      Similar to `tf.py_func` and `tf.py_function` but it doesn't require defining
      the inputs or the dtypes of the outputs a priori.
  
      In Eager mode it would convert the tf.Tensors to np.arrays before passing to
      `func` and then convert back the outputs from np.arrays to tf.Tensors.
  
      In Graph mode it would create different tf.py_function for each combination
      of dtype of the inputs and cache them for reuse.
  
      NOTE: In Graph mode: if `output_dtypes` is not provided then `func` would
      be called with `np.ones()` to infer the output dtypes, and therefore `func`
      should be stateless.
  
      ```python
      Instead of doing:
  
      def sum(x):
        return np.sum(x)
      inputs = tf.constant([3, 4])
      outputs = tf.py_function(sum, inputs, Tout=[tf.int64])
  
      inputs = tf.constant([3., 4.])
      outputs = tf.py_function(sum, inputs, Tout=[tf.float32])
  
      Do:
      @eager_utils.np_function
      def sum(x):
        return np.sum(x)
  
      inputs = tf.constant([3, 4])
      outputs = sum(inputs)  # Infers that Tout is tf.int64
  
      inputs = tf.constant([3., 4.])
      outputs = sum(inputs)  # Infers that Tout is tf.float32
  
      # Output dtype is always float32 for valid input dtypes.
      @eager_utils.np_function(output_dtypes=np.float32)
      def mean(x):
        return np.mean(x)
  
      # Output dtype depends on the input dtype.
      @eager_utils.np_function(output_dtypes=lambda x: (x, x))
      def repeat(x):
        return x, x
  
      with context.graph_mode():
        outputs = sum(tf.constant([3, 4]))
        outputs2 = sum(tf.constant([3., 4.]))
        sess.run(outputs) # np.array(7)
        sess.run(outputs2) # np.array(7.)
  
      with context.eager_mode():
        inputs = tf.constant([3, 4])
        outputs = sum(tf.constant([3, 4])) # tf.Tensor([7])
        outputs = sum(tf.constant([3., 4.])) # tf.Tensor([7.])
  
      ```
      Args:
        func: A numpy function, that takes numpy arrays as inputs and return numpy
          arrays as outputs.
        output_dtypes: Optional list of dtypes or a function that maps input dtypes
          to output dtypes. Examples: output_dtypes=[tf.float32],
          output_dtypes=lambda x: x (outputs have the same dtype as inputs). If it
          is not provided in Graph mode the `func` would be called to infer the
          output dtypes.
  
      Returns:
        A wrapped function that can be used with TF code.
      """
  
      def decorated(func):
          """Decorated func."""
          dtype_map = {}
  
          def wrapper(*args, **kwargs):
              """Wrapper to add nested input and outputs support."""
              func_with_kwargs = functools.partial(func, **kwargs)
  
              def func_flat_outputs(*args):
                  return tf.nest.flatten(func_with_kwargs(*args))
  
              def compute_output_dtypes(*args):
                  """Calls the func to compute output dtypes."""
                  result = func(*args, **kwargs)
                  return tf.nest.map_structure(
                      lambda x: x.dtype, result
                  )
  
              if tf.executing_eagerly():
                  result = func_with_kwargs(
                      *tf.nest.map_structure(lambda x: x.numpy(), args)
                  )
                  convert = (
                      lambda x: x
                      if x is None
                      else tf.convert_to_tensor(value=x)
                  )
                  return tf.nest.map_structure(convert, result)
              else:
                  input_dtypes = tuple(
                      [x.dtype for x in tf.nest.flatten(args)]
                  )
                  if input_dtypes not in dtype_map:
                      if output_dtypes is None:
                          dummy_args = tf.nest.map_structure(
                              lambda x: np.ones(
                                  x.shape, x.dtype.as_numpy_dtype
                              ),
                              args,
                          )
                          dtype_map[
                              input_dtypes
                          ] = compute_output_dtypes(*dummy_args)
                      elif isinstance(output_dtypes, (list, tuple)):
                          # output_dtypes define the output dtypes.
                          dtype_map[input_dtypes] = output_dtypes
                      else:
                          try:
                              # See if output_dtypes define the output dtype directly.
                              tf.as_dtype(output_dtypes)
                              dtype_map[input_dtypes] = output_dtypes
                          except TypeError:
                              if callable(output_dtypes):
                                  # output_dtypes is mapping from input_dtypes to output_dtypes.
                                  dtype_map[
                                      input_dtypes
                                  ] = output_dtypes(*input_dtypes)
                              else:
                                  raise ValueError(
                                      "output_dtypes not a list of dtypes or a callable."
                                  )
  
              flat_output_dtypes = tf.nest.flatten(
                  dtype_map[input_dtypes]
              )
              flat_outputs = tf.py_function(
                  func_flat_outputs, inp=args, Tout=flat_output_dtypes
              )
              return tf.nest.pack_sequence_as(
                  dtype_map[input_dtypes], flat_outputs
              )
  
          return tf_decorator.make_decorator(func, wrapper)
  
      # This code path is for the `foo = np_function(foo, ...)` use case
      if func is not None:
          return decorated(func)
  
      # This code path is for the decorator
      # @np_function(...)
      # def foo(...):
      return decorated
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import copy
  import functools
  import inspect
  from absl import logging
  import numpy
  import six
  import tensorflow
  from tf_agents.utils import common
  from tensorflow.python.util import tf_decorator
label_desc: Use image node_id 10 for calling a global function with example usage: np_function(func, output_dtypes) and returns: decorated, tf_decorator, decorated, tf, tf, tf, tf