tensorflow_probability/python/internal/prefer_static.py
def _prefer_static(original_fn, static_fn):
  """Wraps original_fn, preferring to call static_fn when inputs are static."""
  original_spec = tf_inspect.getfullargspec(original_fn)
  static_spec = tf_inspect.getfullargspec(static_fn)
  if original_spec != static_spec:
    raise ValueError(
        'Arg specs do not match: original={}, static={}, fn={}'.format(
            original_spec, static_spec, original_fn))
  @decorator.decorator
  def wrap(wrapped_fn, *args, **kwargs):
    del wrapped_fn
    [args_, kwargs_], all_static = _maybe_get_static_args([args, kwargs])
    if all_static:
      return static_fn(*args_, **kwargs_)
    return original_fn(*args, **kwargs)
  return wrap(original_fn)
def _prefer_static(original_fn, static_fn):
  """Wraps original_fn, preferring to call static_fn when inputs are static."""
  original_spec = tf_inspect.getfullargspec(original_fn)
  static_spec = tf_inspect.getfullargspec(static_fn)
  if original_spec != static_spec:
    raise ValueError(
        'Arg specs do not match: original={}, static={}, fn={}'.format(
            original_spec, static_spec, original_fn))
  @decorator.decorator
  def wrap(wrapped_fn, *args, **kwargs):
    del wrapped_fn
    [args_, kwargs_], all_static = _maybe_get_static_args([args, kwargs])
    if all_static:
      return static_fn(*args_, **kwargs_)
    return original_fn(*args, **kwargs)
  return wrap(original_fn)
def _prefer_static ( original_fn , static_fn ) : original_spec = tf_inspect . getfullargspec ( original_fn ) static_spec = tf_inspect . getfullargspec ( static_fn ) if original_spec != static_spec : raise ValueError ( 'Arg specs do not match: original={}, static={}, fn={}' . format ( original_spec , static_spec , original_fn ) ) @ decorator . decorator def wrap ( wrapped_fn , * args , * * kwargs ) : del wrapped_fn [ args_ , kwargs_ ] , all_static = _maybe_get_static_args ( [ args , kwargs ] ) if all_static : return static_fn ( * args_ , * * kwargs_ ) return original_fn ( * args , * * kwargs ) return wrap ( original_fn )
Wraps original_fn, preferring to call static_fn when inputs are static.
Wraps original_fn preferring to call static_fn when inputs are static .