tensorflow_probability/python/mcmc/langevin.py
def _maybe_call_volatility_fn_and_grads(volatility_fn,
                                        state,
                                        volatility_fn_results=None,
                                        grads_volatility_fn=None,
                                        sample_shape=None,
                                        parallel_iterations=10):
  """Helper which computes `volatility_fn` results and grads, if needed."""
  state_parts = list(state) if mcmc_util.is_list_like(state) else [state]
  needs_volatility_fn_gradients = grads_volatility_fn is None

  # Convert `volatility_fn_results` to a list
  if volatility_fn_results is None:
    volatility_fn_results = volatility_fn(*state_parts)

  volatility_fn_results = (list(volatility_fn_results)
                           if mcmc_util.is_list_like(volatility_fn_results)
                           else [volatility_fn_results])
  if len(volatility_fn_results) == 1:
    volatility_fn_results *= len(state_parts)
  if len(state_parts) != len(volatility_fn_results):
    raise ValueError('`volatility_fn` should return a tensor or a list '
                     'of the same length as `current_state`.')

  # The shape of 'volatility_parts' needs to have the number of chains as a
  # leading dimension. For determinism we broadcast 'volatility_parts' to the
  # shape of `state_parts` since each dimension of `state_parts` could have a
  # different volatility value.

  volatility_fn_results = _maybe_broadcast_volatility(volatility_fn_results,
                                                      state_parts)
  if grads_volatility_fn is None:
    [
        _,
        grads_volatility_fn,
    ] = diag_jacobian(
        xs=state_parts,
        ys=volatility_fn_results,
        sample_shape=sample_shape,
        parallel_iterations=parallel_iterations,
        fn=volatility_fn)

  # Compute gradient of `volatility_parts**2`
  if needs_volatility_fn_gradients:
    grads_volatility_fn = [
        2. * g * volatility if g is not None else tf.zeros_like(
            fn_arg, dtype=fn_arg.dtype.base_dtype)
        for g, volatility, fn_arg in zip(
            grads_volatility_fn, volatility_fn_results, state_parts)
    ]

  return volatility_fn_results, grads_volatility_fn
def _maybe_call_volatility_fn_and_grads(volatility_fn,
                                        state,
                                        volatility_fn_results=None,
                                        grads_volatility_fn=None,
                                        sample_shape=None,
                                        parallel_iterations=10):
  state_parts = list(state) if mcmc_util.is_list_like(state) else [state]
  needs_volatility_fn_gradients = grads_volatility_fn is None

  # Convert `volatility_fn_results` to a list
  if volatility_fn_results is None:
    volatility_fn_results = volatility_fn(*state_parts)

  volatility_fn_results = (list(volatility_fn_results)
                           if mcmc_util.is_list_like(volatility_fn_results)
                           else [volatility_fn_results])
  if len(volatility_fn_results) == 1:
    volatility_fn_results *= len(state_parts)
  if len(state_parts) != len(volatility_fn_results):
    raise ValueError('`volatility_fn` should return a tensor or a list '
                     'of the same length as `current_state`.')

  # The shape of 'volatility_parts' needs to have the number of chains as a
  # leading dimension. For determinism we broadcast 'volatility_parts' to the
  # shape of `state_parts` since each dimension of `state_parts` could have a
  # different volatility value.

  volatility_fn_results = _maybe_broadcast_volatility(volatility_fn_results,
                                                      state_parts)
  if grads_volatility_fn is None:
    [
        _,
        grads_volatility_fn,
    ] = diag_jacobian(
        xs=state_parts,
        ys=volatility_fn_results,
        sample_shape=sample_shape,
        parallel_iterations=parallel_iterations,
        fn=volatility_fn)

  # Compute gradient of `volatility_parts**2`
  if needs_volatility_fn_gradients:
    grads_volatility_fn = [
        2. * g * volatility if g is not None else tf.zeros_like(
            fn_arg, dtype=fn_arg.dtype.base_dtype)
        for g, volatility, fn_arg in zip(
            grads_volatility_fn, volatility_fn_results, state_parts)
    ]

  return volatility_fn_results, grads_volatility_fn
def _maybe_call_volatility_fn_and_grads ( volatility_fn , state , volatility_fn_results = None , grads_volatility_fn = None , sample_shape = None , parallel_iterations = 10 ) : state_parts = list ( state ) if mcmc_util . is_list_like ( state ) else [ state ] needs_volatility_fn_gradients = grads_volatility_fn is None # Convert `volatility_fn_results` to a list if volatility_fn_results is None : volatility_fn_results = volatility_fn ( * state_parts ) volatility_fn_results = ( list ( volatility_fn_results ) if mcmc_util . is_list_like ( volatility_fn_results ) else [ volatility_fn_results ] ) if len ( volatility_fn_results ) == 1 : volatility_fn_results *= len ( state_parts ) if len ( state_parts ) != len ( volatility_fn_results ) : raise ValueError ( '`volatility_fn` should return a tensor or a list ' 'of the same length as `current_state`.' ) # The shape of 'volatility_parts' needs to have the number of chains as a # leading dimension. For determinism we broadcast 'volatility_parts' to the # shape of `state_parts` since each dimension of `state_parts` could have a # different volatility value. volatility_fn_results = _maybe_broadcast_volatility ( volatility_fn_results , state_parts ) if grads_volatility_fn is None : [ _ , grads_volatility_fn , ] = diag_jacobian ( xs = state_parts , ys = volatility_fn_results , sample_shape = sample_shape , parallel_iterations = parallel_iterations , fn = volatility_fn ) # Compute gradient of `volatility_parts**2` if needs_volatility_fn_gradients : grads_volatility_fn = [ 2. * g * volatility if g is not None else tf . zeros_like ( fn_arg , dtype = fn_arg . dtype . base_dtype ) for g , volatility , fn_arg in zip ( grads_volatility_fn , volatility_fn_results , state_parts ) ] return volatility_fn_results , grads_volatility_fn
Helper which computes `volatility_fn` results and grads, if needed.
Helper which computes volatility_fn results and grads if needed .