code: def tf_abstract_eval(f):
      """Returns a function that evaluates `f` given input shapes and dtypes.
  
      It transforms function `f` to a function that performs the same computation as
      `f` but only on shapes and dtypes (a.k.a. shape inference).
  
      Args:
        f: the function to be transformed.
  
      Returns:
        A function whose input arguments can be either the same as `f`'s or only
        their shapes/dtypes represented by `ShapeDtype`, and whose return values are
        `ShapeDtype`s with the same nested structure as `f`'s return values.
      """
      f_shape = tf_np_extensions.eval_on_shapes(f)
  
      def from_shape_type(x):
          if isinstance(x, ShapeDtype):
              return tf.TensorSpec(x.shape, x.dtype)
          else:
              return x
  
      def to_shape_type(x):  # pylint: disable=missing-docstring
          # TODO(wangpeng): handle partial output shapes using `tf.shape`.
          def to_numpy_shape(s):
              if s.is_fully_defined():
                  return tuple(s.as_list())
              else:
                  raise ValueError(
                      "The output shapes (%s) of the dry-run'ed function are"
                      " not fully defined." % s
                  )
  
          def to_numpy_dtype(t):
              return np.dtype(t.as_numpy_dtype)
  
          if isinstance(x, tf.TensorSpec):
              return ShapeDtype(
                  to_numpy_shape(x.shape), to_numpy_dtype(x.dtype)
              )
          else:
              return x
  
      def f_return(*args):
          args = tf.nest.map_structure(from_shape_type, args)
          res = f_shape(*args)
          return tf.nest.map_structure(to_shape_type, res)
  
      return f_return
  
file_path: trax/trax/fastmath/tf.py
imports: import numpy
  import tensorflow.compat.v2
  from trax.shapes import ShapeDtype
  from trax.tf_numpy import extensions
  from trax.tf_numpy import numpy
label_desc: Use image node_id 1 for calling a global function with example usage: tf_abstract_eval(f) and returns: f_return, tf, tf, x, np, ShapeDtype, x, tuple