code: def jax_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.
      """
  
      def shape_fun(*args, **kwargs):
          jax_shapes = jax.eval_shape(f, *args, **kwargs)
          return tnp.nested_map(signature, jax_shapes)
  
      return shape_fun
  
imports: import functools
  import jax
  from jax import lax
  from jax import random
  import jax.numpy
  import jax.scipy.special
  import numpy
  import tensorflow
  import tensorflow_datasets
  from trax.fastmath import numpy
  from trax.shapes import signature
label_desc: Use image node_id 7 for calling a global function with example usage: jax_abstract_eval(f) and returns: shape_fun, tnp