(f, static_argnums=(), allow_static_outputs=False)
return f_return
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.
    static_argnums: see documentation of `jit`.
    allow_static_outputs: whether to allow non-array outputs. If True, non-array
      outputs (e.g. Python integers) will be returned as-is; otherwise, they
      will be converted to ndarrays, and then specs of those ndarrays will be
      returned.

  Returns:
    A function whose input arguments can be either the same as `f`'s or only
    their shapes/dtypes represented by `tf.TensorSpec`, and whose return values
    are `tf.TensorSpec`s with the same nested structure as `f`'s return
    values. If `allow_static_outputs` is True, when `f` returns some non-array
    outputs (e.g. Python integers), the converted function will return them
    as-is instead of returning `tf.TensorSpec`s for them.
Returns a function that evaluates `f` given input shapes and dtypes.
Returns a function that evaluates f given input shapes and dtypes .
def eval_on_shapes(f, static_argnums=(), allow_static_outputs=False):
  """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.
    static_argnums: see documentation of `jit`.
    allow_static_outputs: whether to allow non-array outputs. If True, non-array
      outputs (e.g. Python integers) will be returned as-is; otherwise, they
      will be converted to ndarrays, and then specs of those ndarrays will be
      returned.

  Returns:
    A function whose input arguments can be either the same as `f`'s or only
    their shapes/dtypes represented by `tf.TensorSpec`, and whose return values
    are `tf.TensorSpec`s with the same nested structure as `f`'s return
    values. If `allow_static_outputs` is True, when `f` returns some non-array
    outputs (e.g. Python integers), the converted function will return them
    as-is instead of returning `tf.TensorSpec`s for them.
  """
  def abstractify(args):
    def _abstractify(x):
      x = _canonicalize_jit_arg(x)
      if isinstance(x, (tf.Tensor, tf_np.ndarray)):
        return tf.TensorSpec(x.shape, x.dtype)
      else:
        return x
    new_args = []
    for i, arg in enumerate(args):
      if i in static_argnums:
        new_args.append(arg)
      else:
        new_args.append(tf.nest.map_structure(_abstractify, arg))
    return new_args

  if allow_static_outputs:
    # When `tf_f` below is called (via get_concrete_function) with the same
    # arugments (after abstraction), the Python function `f` won't be run, so we
    # need this python_outputs_map to retrieve the Python outputs we've seen
    # before that correspond the arguments.
    python_outputs_map = {}
    def recorder(args, kwargs, res):
      # Since the get_concrete_function below only uses positional args, we also
      # only positional args here.
      del args, kwargs
      def is_tensor_like(x):
        if hasattr(x, "_type_spec"):
          return True  # x is a CompositeTensor
        return isinstance(x, (tf_np.ndarray, tf.Tensor))
      py_values = tf.nest.map_structure(
          lambda x: None if is_tensor_like(x) else x,
          res)
      key = id(tf.compat.v1.get_default_graph())
      python_outputs_map[key] = py_values
      # Set non-tensor outputs to None to avoid tf.function calling
      # tf.convert_to_tensor on them.
      res = tf.nest.map_structure(
          lambda x: None if not is_tensor_like(x) else x,
          res)
      return res
    f = _record_result_type(recorder, f)

  # TODO(wangpeng): tf.function could add a knob to turn off materializing the
  #   graph, so that we don't waste computation and memory when we just want
  #   shape inference.
  tf_f = jit(f, static_argnums=static_argnums).tf_function

  # pylint: disable=missing-docstring
  def f_return(*args):
    def to_tensor_spec(x):
      if isinstance(x, tf.Tensor):
        return tf.TensorSpec(x.shape, x.dtype)
      else:
        return x

    new_args = abstractify(args)
    cfun = tf_f.get_concrete_function(*new_args)
    res = cfun.structured_outputs
    res = tf.nest.map_structure(to_tensor_spec, res)

    if allow_static_outputs:
      key = id(cfun.graph)
      py_values = python_outputs_map[key]
      # We can also call tf.get_static_value on structured_outputs to retrieve
      # the Python values, but since we'll need to use python_outputs_map to
      # record "which outputs are static?" anyway, we choose to directly store
      # the Python values in python_outputs_map.
      res = tf.nest.map_structure(
          lambda x, python_value: x if python_value is None else python_value,
          res, py_values)

    return res

  # Provides access to `tf_f` for testing purpose.
  f_return._tf_function = tf_f  # pylint: disable=protected-access
  return f_return
def eval_on_shapes ( f , static_argnums = ( ) , allow_static_outputs = False ) : def abstractify ( args ) : def _abstractify ( x ) : x = _canonicalize_jit_arg ( x ) if isinstance ( x , ( tf . Tensor , tf_np . ndarray ) ) : return tf . TensorSpec ( x . shape , x . dtype ) else : return x new_args = [ ] for i , arg in enumerate ( args ) : if i in static_argnums : new_args . append ( arg ) else : new_args . append ( tf . nest . map_structure ( _abstractify , arg ) ) return new_args if allow_static_outputs : # When `tf_f` below is called (via get_concrete_function) with the same # arugments (after abstraction), the Python function `f` won't be run, so we # need this python_outputs_map to retrieve the Python outputs we've seen # before that correspond the arguments. python_outputs_map = { } def recorder ( args , kwargs , res ) : # Since the get_concrete_function below only uses positional args, we also # only positional args here. del args , kwargs def is_tensor_like ( x ) : if hasattr ( x , "_type_spec" ) : return True # x is a CompositeTensor return isinstance ( x , ( tf_np . ndarray , tf . Tensor ) ) py_values = tf . nest . map_structure ( lambda x : None if is_tensor_like ( x ) else x , res ) key = id ( tf . compat . v1 . get_default_graph ( ) ) python_outputs_map [ key ] = py_values # Set non-tensor outputs to None to avoid tf.function calling # tf.convert_to_tensor on them. res = tf . nest . map_structure ( lambda x : None if not is_tensor_like ( x ) else x , res ) return res f = _record_result_type ( recorder , f ) # TODO(wangpeng): tf.function could add a knob to turn off materializing the #   graph, so that we don't waste computation and memory when we just want #   shape inference. tf_f = jit ( f , static_argnums = static_argnums ) . tf_function # pylint: disable=missing-docstring def f_return ( * args ) : def to_tensor_spec ( x ) : if isinstance ( x , tf . Tensor ) : return tf . TensorSpec ( x . shape , x . dtype ) else : return x new_args = abstractify ( args ) cfun = tf_f . get_concrete_function ( * new_args ) res = cfun . structured_outputs res = tf . nest . map_structure ( to_tensor_spec , res ) if allow_static_outputs : key = id ( cfun . graph ) py_values = python_outputs_map [ key ] # We can also call tf.get_static_value on structured_outputs to retrieve # the Python values, but since we'll need to use python_outputs_map to # record "which outputs are static?" anyway, we choose to directly store # the Python values in python_outputs_map. res = tf . nest . map_structure ( lambda x , python_value : x if python_value is None else python_value , res , py_values ) return res # Provides access to `tf_f` for testing purpose. f_return . _tf_function = tf_f # pylint: disable=protected-access return f_return