tensorflow_probability/python/mcmc/internal/util.py
def _value_and_gradients(fn, fn_arg_list, result=None, grads=None, name=None):
  """Helper to `maybe_call_fn_and_grads`."""
  with tf.compat.v1.name_scope(name, 'value_and_gradients',
                               [fn_arg_list, result, grads]):

    def _convert_to_tensor(x, name):
      ctt = lambda x_: x_ if x_ is None else tf.convert_to_tensor(
          value=x_, name=name)
      return [ctt(x_) for x_ in x] if is_list_like(x) else ctt(x)

    fn_arg_list = (list(fn_arg_list) if is_list_like(fn_arg_list)
                   else [fn_arg_list])
    fn_arg_list = _convert_to_tensor(fn_arg_list, 'fn_arg')

    if result is None:
      result = fn(*fn_arg_list)
      if grads is None and tf.executing_eagerly():
        # Ensure we disable bijector cacheing in eager mode.
        # TODO(b/72831017): Remove this once bijector cacheing is fixed for
        # eager mode.
        fn_arg_list = [0 + x for x in fn_arg_list]

    result = _convert_to_tensor(result, 'fn_result')

    if grads is not None:
      grads = _convert_to_tensor(grads, 'fn_grad')
      return result, grads

    if is_list_like(result) and len(result) == len(fn_arg_list):
      # Compute the block diagonal of Jacobian.
      # TODO(b/79158574): Guard this calculation by an arg which explicitly
      # requests block diagonal Jacobian calculation.
      def fn_slice(i):
        """Needed to prevent `cell-var-from-loop` pylint warning."""
        return lambda x: fn(*(fn_arg_list[:i] + [x] + fn_arg_list[i+1:]))
      grads = [
          tfp_math_value_and_gradients(fn_slice(i), fn_arg_list[i])[1]
          for i in range(len(result))
      ]
    else:
      _, grads = tfp_math_value_and_gradients(fn, fn_arg_list)

    return result, grads
def _value_and_gradients(fn, fn_arg_list, result=None, grads=None, name=None):
  with tf.compat.v1.name_scope(name, 'value_and_gradients',
                               [fn_arg_list, result, grads]):

    def _convert_to_tensor(x, name):
      ctt = lambda x_: x_ if x_ is None else tf.convert_to_tensor(
          value=x_, name=name)
      return [ctt(x_) for x_ in x] if is_list_like(x) else ctt(x)

    fn_arg_list = (list(fn_arg_list) if is_list_like(fn_arg_list)
                   else [fn_arg_list])
    fn_arg_list = _convert_to_tensor(fn_arg_list, 'fn_arg')

    if result is None:
      result = fn(*fn_arg_list)
      if grads is None and tf.executing_eagerly():
        # Ensure we disable bijector cacheing in eager mode.
        # TODO(b/72831017): Remove this once bijector cacheing is fixed for
        # eager mode.
        fn_arg_list = [0 + x for x in fn_arg_list]

    result = _convert_to_tensor(result, 'fn_result')

    if grads is not None:
      grads = _convert_to_tensor(grads, 'fn_grad')
      return result, grads

    if is_list_like(result) and len(result) == len(fn_arg_list):
      # Compute the block diagonal of Jacobian.
      # TODO(b/79158574): Guard this calculation by an arg which explicitly
      # requests block diagonal Jacobian calculation.
      def fn_slice(i):
        """Needed to prevent `cell-var-from-loop` pylint warning."""
        return lambda x: fn(*(fn_arg_list[:i] + [x] + fn_arg_list[i+1:]))
      grads = [
          tfp_math_value_and_gradients(fn_slice(i), fn_arg_list[i])[1]
          for i in range(len(result))
      ]
    else:
      _, grads = tfp_math_value_and_gradients(fn, fn_arg_list)

    return result, grads
def _value_and_gradients ( fn , fn_arg_list , result = None , grads = None , name = None ) : with tf . compat . v1 . name_scope ( name , 'value_and_gradients' , [ fn_arg_list , result , grads ] ) : def _convert_to_tensor ( x , name ) : ctt = lambda x_ : x_ if x_ is None else tf . convert_to_tensor ( value = x_ , name = name ) return [ ctt ( x_ ) for x_ in x ] if is_list_like ( x ) else ctt ( x ) fn_arg_list = ( list ( fn_arg_list ) if is_list_like ( fn_arg_list ) else [ fn_arg_list ] ) fn_arg_list = _convert_to_tensor ( fn_arg_list , 'fn_arg' ) if result is None : result = fn ( * fn_arg_list ) if grads is None and tf . executing_eagerly ( ) : # Ensure we disable bijector cacheing in eager mode. # TODO(b/72831017): Remove this once bijector cacheing is fixed for # eager mode. fn_arg_list = [ 0 + x for x in fn_arg_list ] result = _convert_to_tensor ( result , 'fn_result' ) if grads is not None : grads = _convert_to_tensor ( grads , 'fn_grad' ) return result , grads if is_list_like ( result ) and len ( result ) == len ( fn_arg_list ) : # Compute the block diagonal of Jacobian. # TODO(b/79158574): Guard this calculation by an arg which explicitly # requests block diagonal Jacobian calculation. def fn_slice ( i ) : """Needed to prevent `cell-var-from-loop` pylint warning.""" return lambda x : fn ( * ( fn_arg_list [ : i ] + [ x ] + fn_arg_list [ i + 1 : ] ) ) grads = [ tfp_math_value_and_gradients ( fn_slice ( i ) , fn_arg_list [ i ] ) [ 1 ] for i in range ( len ( result ) ) ] else : _ , grads = tfp_math_value_and_gradients ( fn , fn_arg_list ) return result , grads
Helper to `maybe_call_fn_and_grads`.
Helper to maybe_call_fn_and_grads .