experimental/fun_mcmc/fun_mcmc_lib.py
def hamiltonian_monte_carlo(
    hmc_state: HamiltonianMonteCarloState,
    target_log_prob_fn: PotentialFn,
    step_size: Any,
    num_leapfrog_steps: IntTensor,
    momentum: State = None,
    kinetic_energy_fn: PotentialFn = None,
    momentum_sample_fn: MomentumSampleFn = None,
    leapfrog_trace_fn: Callable[[LeapFrogStepState, LeapFrogStepExtras],
                                TensorNest] = lambda *args: (),
    seed=None,
) -> Tuple[HamiltonianMonteCarloState, HamiltonianMonteCarloExtra]:
  """Hamiltonian Monte Carlo `TransitionOperator`.

  #### Example

  ```python
  step_size = 0.2
  num_steps = 2000
  num_leapfrog_steps = 10
  state = tf.ones([16, 2])

  base_mean = [1., 0]
  base_cov = [[1, 0.5], [0.5, 1]]

  bijector = tfb.Softplus()
  base_dist = tfd.MultivariateNormalFullCovariance(
      loc=base_mean, covariance_matrix=base_cov)
  target_dist = bijector(base_dist)

  def orig_target_log_prob_fn(x):
    return target_dist.log_prob(x), ()

  target_log_prob_fn, state = fun_mcmc.transform_log_prob_fn(
      orig_target_log_prob_fn, bijector, state)

  kernel = tf.function(lambda state: fun_mcmc.hamiltonian_monte_carlo(
      state,
      step_size=step_size,
      num_leapfrog_steps=num_leapfrog_steps,
      target_log_prob_fn=target_log_prob_fn,
      seed=tfp_test_util.test_seed()))

  _, chain = fun_mcmc.trace(
      state=fun_mcmc.HamiltonianMonteCarloState(
          state=state,
          state_grads=None,
          target_log_prob=None,
          state_extra=None),
      fn=kernel,
      num_steps=num_steps,
      trace_fn=lambda state, extra: state.state_extra[0])
  ```

  Args:
    hmc_state: HamiltonianMonteCarloState.
    target_log_prob_fn: Target log prob fn.
    step_size: Step size, structure broadcastable to the `target_log_prob_fn`
      state.
    num_leapfrog_steps: Number of leapfrog steps to take.
    momentum: Initial momentum, passed to `momentum_sample_fn`. Default: zeroes.
    kinetic_energy_fn: Kinetic energy function.
    momentum_sample_fn: Sampler for the momentum.
    leapfrog_trace_fn: Trace function for the leapfrog integrator.
    seed: For reproducibility.

  Returns:
    hmc_state: HamiltonianMonteCarloState
    hmc_extra: HamiltonianMonteCarloExtra
  """
  state = hmc_state.state
  state_grads = hmc_state.state_grads
  target_log_prob = hmc_state.target_log_prob
  state_extra = hmc_state.state_extra

  if kinetic_energy_fn is None:

    # pylint: disable=function-redefined
    def kinetic_energy_fn(*momentum):
      return tf.add_n([
          tf.reduce_sum(input_tensor=tf.square(x), axis=-1) / 2.
          for x in tf.nest.flatten(momentum)
      ]), ()

  if momentum_sample_fn is None:

    # pylint: disable=function-redefined
    def momentum_sample_fn(*momentum):
      ret = tf.nest.map_structure(
          lambda x: tf.random.normal(tf.shape(input=x), dtype=x.dtype),
          momentum)
      if len(ret) == 1:
        return ret[0]
      else:
        return ret

  if momentum is None:
    momentum = call_fn(momentum_sample_fn,
                       tf.nest.map_structure(tf.zeros_like, state))
  if target_log_prob is None:
    target_log_prob, state_extra, state_grads = call_and_grads(
        target_log_prob_fn, state)

  kinetic_energy, _ = call_fn(kinetic_energy_fn, momentum)
  current_energy = -target_log_prob + kinetic_energy
  current_state = HamiltonianMonteCarloState(
      state=state,
      state_grads=state_grads,
      state_extra=state_extra,
      target_log_prob=target_log_prob)

  def leapfrog_wrapper(leapfrog_state, target_log_prob, state_extra):
    """Leapfrog wrapper that tracks extra state."""
    del target_log_prob
    del state_extra

    leapfrog_state, leapfrog_extra = leapfrog_step(
        leapfrog_state,
        step_size=step_size,
        target_log_prob_fn=target_log_prob_fn,
        kinetic_energy_fn=kinetic_energy_fn)

    return [
        leapfrog_state, leapfrog_extra.target_log_prob,
        leapfrog_extra.state_extra
    ], leapfrog_extra

  def leapfrog_trace_wrapper_fn(args, leapfrog_extra):
    return leapfrog_trace_fn(args[0], leapfrog_extra)

  leapfrog_wrapper_state = (LeapFrogStepState(state, state_grads, momentum),
                            target_log_prob, state_extra)

  [[leapfrog_state, target_log_prob, state_extra], _], leapfrog_trace = trace(
      leapfrog_wrapper_state,
      leapfrog_wrapper,
      num_leapfrog_steps,
      trace_fn=leapfrog_trace_wrapper_fn)

  kinetic_energy, _ = call_fn(kinetic_energy_fn, leapfrog_state.momentum)
  proposed_energy = -target_log_prob + kinetic_energy
  proposed_state = HamiltonianMonteCarloState(
      state=leapfrog_state.state,
      state_grads=leapfrog_state.state_grads,
      target_log_prob=target_log_prob,
      state_extra=state_extra)

  energy_change = proposed_energy - current_energy
  hmc_state, is_accepted, _ = metropolis_hastings_step(
      current_state, proposed_state, energy_change, seed=seed)

  hmc_state = hmc_state  # type: HamiltonianMonteCarloState
  return hmc_state, HamiltonianMonteCarloExtra(
      is_accepted=is_accepted,
      proposed_hmc_state=proposed_state,
      log_accept_ratio=-energy_change,
      leapfrog_trace=leapfrog_trace)
def hamiltonian_monte_carlo(
    hmc_state: HamiltonianMonteCarloState,
    target_log_prob_fn: PotentialFn,
    step_size: Any,
    num_leapfrog_steps: IntTensor,
    momentum: State = None,
    kinetic_energy_fn: PotentialFn = None,
    momentum_sample_fn: MomentumSampleFn = None,
    leapfrog_trace_fn: Callable[[LeapFrogStepState, LeapFrogStepExtras],
                                TensorNest] = lambda *args: (),
    seed=None,
) -> Tuple[HamiltonianMonteCarloState, HamiltonianMonteCarloExtra]:
  """Hamiltonian Monte Carlo `TransitionOperator`.

  #### Example

  ```python
  step_size = 0.2
  num_steps = 2000
  num_leapfrog_steps = 10
  state = tf.ones([16, 2])

  base_mean = [1., 0]
  base_cov = [[1, 0.5], [0.5, 1]]

  bijector = tfb.Softplus()
  base_dist = tfd.MultivariateNormalFullCovariance(
      loc=base_mean, covariance_matrix=base_cov)
  target_dist = bijector(base_dist)

  def orig_target_log_prob_fn(x):
    return target_dist.log_prob(x), ()

  target_log_prob_fn, state = fun_mcmc.transform_log_prob_fn(
      orig_target_log_prob_fn, bijector, state)

  kernel = tf.function(lambda state: fun_mcmc.hamiltonian_monte_carlo(
      state,
      step_size=step_size,
      num_leapfrog_steps=num_leapfrog_steps,
      target_log_prob_fn=target_log_prob_fn,
      seed=tfp_test_util.test_seed()))

  _, chain = fun_mcmc.trace(
      state=fun_mcmc.HamiltonianMonteCarloState(
          state=state,
          state_grads=None,
          target_log_prob=None,
          state_extra=None),
      fn=kernel,
      num_steps=num_steps,
      trace_fn=lambda state, extra: state.state_extra[0])
  ```

  Args:
    hmc_state: HamiltonianMonteCarloState.
    target_log_prob_fn: Target log prob fn.
    step_size: Step size, structure broadcastable to the `target_log_prob_fn`
      state.
    num_leapfrog_steps: Number of leapfrog steps to take.
    momentum: Initial momentum, passed to `momentum_sample_fn`. Default: zeroes.
    kinetic_energy_fn: Kinetic energy function.
    momentum_sample_fn: Sampler for the momentum.
    leapfrog_trace_fn: Trace function for the leapfrog integrator.
    seed: For reproducibility.

  Returns:
    hmc_state: HamiltonianMonteCarloState
    hmc_extra: HamiltonianMonteCarloExtra
  """
  state = hmc_state.state
  state_grads = hmc_state.state_grads
  target_log_prob = hmc_state.target_log_prob
  state_extra = hmc_state.state_extra

  if kinetic_energy_fn is None:

    # pylint: disable=function-redefined
    def kinetic_energy_fn(*momentum):
      return tf.add_n([
          tf.reduce_sum(input_tensor=tf.square(x), axis=-1) / 2.
          for x in tf.nest.flatten(momentum)
      ]), ()

  if momentum_sample_fn is None:

    # pylint: disable=function-redefined
    def momentum_sample_fn(*momentum):
      ret = tf.nest.map_structure(
          lambda x: tf.random.normal(tf.shape(input=x), dtype=x.dtype),
          momentum)
      if len(ret) == 1:
        return ret[0]
      else:
        return ret

  if momentum is None:
    momentum = call_fn(momentum_sample_fn,
                       tf.nest.map_structure(tf.zeros_like, state))
  if target_log_prob is None:
    target_log_prob, state_extra, state_grads = call_and_grads(
        target_log_prob_fn, state)

  kinetic_energy, _ = call_fn(kinetic_energy_fn, momentum)
  current_energy = -target_log_prob + kinetic_energy
  current_state = HamiltonianMonteCarloState(
      state=state,
      state_grads=state_grads,
      state_extra=state_extra,
      target_log_prob=target_log_prob)

  def leapfrog_wrapper(leapfrog_state, target_log_prob, state_extra):
    """Leapfrog wrapper that tracks extra state."""
    del target_log_prob
    del state_extra

    leapfrog_state, leapfrog_extra = leapfrog_step(
        leapfrog_state,
        step_size=step_size,
        target_log_prob_fn=target_log_prob_fn,
        kinetic_energy_fn=kinetic_energy_fn)

    return [
        leapfrog_state, leapfrog_extra.target_log_prob,
        leapfrog_extra.state_extra
    ], leapfrog_extra

  def leapfrog_trace_wrapper_fn(args, leapfrog_extra):
    return leapfrog_trace_fn(args[0], leapfrog_extra)

  leapfrog_wrapper_state = (LeapFrogStepState(state, state_grads, momentum),
                            target_log_prob, state_extra)

  [[leapfrog_state, target_log_prob, state_extra], _], leapfrog_trace = trace(
      leapfrog_wrapper_state,
      leapfrog_wrapper,
      num_leapfrog_steps,
      trace_fn=leapfrog_trace_wrapper_fn)

  kinetic_energy, _ = call_fn(kinetic_energy_fn, leapfrog_state.momentum)
  proposed_energy = -target_log_prob + kinetic_energy
  proposed_state = HamiltonianMonteCarloState(
      state=leapfrog_state.state,
      state_grads=leapfrog_state.state_grads,
      target_log_prob=target_log_prob,
      state_extra=state_extra)

  energy_change = proposed_energy - current_energy
  hmc_state, is_accepted, _ = metropolis_hastings_step(
      current_state, proposed_state, energy_change, seed=seed)

  hmc_state = hmc_state  # type: HamiltonianMonteCarloState
  return hmc_state, HamiltonianMonteCarloExtra(
      is_accepted=is_accepted,
      proposed_hmc_state=proposed_state,
      log_accept_ratio=-energy_change,
      leapfrog_trace=leapfrog_trace)
def hamiltonian_monte_carlo ( hmc_state : HamiltonianMonteCarloState , target_log_prob_fn : PotentialFn , step_size : Any , num_leapfrog_steps : IntTensor , momentum : State = None , kinetic_energy_fn : PotentialFn = None , momentum_sample_fn : MomentumSampleFn = None , leapfrog_trace_fn : Callable [ [ LeapFrogStepState , LeapFrogStepExtras ] , TensorNest ] = lambda * args : ( ) , seed = None , ) -> Tuple [ HamiltonianMonteCarloState , HamiltonianMonteCarloExtra ] : state = hmc_state . state state_grads = hmc_state . state_grads target_log_prob = hmc_state . target_log_prob state_extra = hmc_state . state_extra if kinetic_energy_fn is None : # pylint: disable=function-redefined def kinetic_energy_fn ( * momentum ) : return tf . add_n ( [ tf . reduce_sum ( input_tensor = tf . square ( x ) , axis = - 1 ) / 2. for x in tf . nest . flatten ( momentum ) ] ) , ( ) if momentum_sample_fn is None : # pylint: disable=function-redefined def momentum_sample_fn ( * momentum ) : ret = tf . nest . map_structure ( lambda x : tf . random . normal ( tf . shape ( input = x ) , dtype = x . dtype ) , momentum ) if len ( ret ) == 1 : return ret [ 0 ] else : return ret if momentum is None : momentum = call_fn ( momentum_sample_fn , tf . nest . map_structure ( tf . zeros_like , state ) ) if target_log_prob is None : target_log_prob , state_extra , state_grads = call_and_grads ( target_log_prob_fn , state ) kinetic_energy , _ = call_fn ( kinetic_energy_fn , momentum ) current_energy = - target_log_prob + kinetic_energy current_state = HamiltonianMonteCarloState ( state = state , state_grads = state_grads , state_extra = state_extra , target_log_prob = target_log_prob ) def leapfrog_wrapper ( leapfrog_state , target_log_prob , state_extra ) : """Leapfrog wrapper that tracks extra state.""" del target_log_prob del state_extra leapfrog_state , leapfrog_extra = leapfrog_step ( leapfrog_state , step_size = step_size , target_log_prob_fn = target_log_prob_fn , kinetic_energy_fn = kinetic_energy_fn ) return [ leapfrog_state , leapfrog_extra . target_log_prob , leapfrog_extra . state_extra ] , leapfrog_extra def leapfrog_trace_wrapper_fn ( args , leapfrog_extra ) : return leapfrog_trace_fn ( args [ 0 ] , leapfrog_extra ) leapfrog_wrapper_state = ( LeapFrogStepState ( state , state_grads , momentum ) , target_log_prob , state_extra ) [ [ leapfrog_state , target_log_prob , state_extra ] , _ ] , leapfrog_trace = trace ( leapfrog_wrapper_state , leapfrog_wrapper , num_leapfrog_steps , trace_fn = leapfrog_trace_wrapper_fn ) kinetic_energy , _ = call_fn ( kinetic_energy_fn , leapfrog_state . momentum ) proposed_energy = - target_log_prob + kinetic_energy proposed_state = HamiltonianMonteCarloState ( state = leapfrog_state . state , state_grads = leapfrog_state . state_grads , target_log_prob = target_log_prob , state_extra = state_extra ) energy_change = proposed_energy - current_energy hmc_state , is_accepted , _ = metropolis_hastings_step ( current_state , proposed_state , energy_change , seed = seed ) hmc_state = hmc_state # type: HamiltonianMonteCarloState return hmc_state , HamiltonianMonteCarloExtra ( is_accepted = is_accepted , proposed_hmc_state = proposed_state , log_accept_ratio = - energy_change , leapfrog_trace = leapfrog_trace )
Hamiltonian Monte Carlo `TransitionOperator`.

  #### Example

  ```python
  step_size = 0.2
  num_steps = 2000
  num_leapfrog_steps = 10
  state = tf.ones([16, 2])

  base_mean = [1., 0]
  base_cov = [[1, 0.5], [0.5, 1]]

  bijector = tfb.Softplus()
  base_dist = tfd.MultivariateNormalFullCovariance(
      loc=base_mean, covariance_matrix=base_cov)
  target_dist = bijector(base_dist)

  def orig_target_log_prob_fn(x):
    return target_dist.log_prob(x), ()

  target_log_prob_fn, state = fun_mcmc.transform_log_prob_fn(
      orig_target_log_prob_fn, bijector, state)

  kernel = tf.function(lambda state: fun_mcmc.hamiltonian_monte_carlo(
      state,
      step_size=step_size,
      num_leapfrog_steps=num_leapfrog_steps,
      target_log_prob_fn=target_log_prob_fn,
      seed=tfp_test_util.test_seed()))

  _, chain = fun_mcmc.trace(
      state=fun_mcmc.HamiltonianMonteCarloState(
          state=state,
          state_grads=None,
          target_log_prob=None,
          state_extra=None),
      fn=kernel,
      num_steps=num_steps,
      trace_fn=lambda state, extra: state.state_extra[0])
  ```

  Args:
    hmc_state: HamiltonianMonteCarloState.
    target_log_prob_fn: Target log prob fn.
    step_size: Step size, structure broadcastable to the `target_log_prob_fn`
      state.
    num_leapfrog_steps: Number of leapfrog steps to take.
    momentum: Initial momentum, passed to `momentum_sample_fn`. Default: zeroes.
    kinetic_energy_fn: Kinetic energy function.
    momentum_sample_fn: Sampler for the momentum.
    leapfrog_trace_fn: Trace function for the leapfrog integrator.
    seed: For reproducibility.

  Returns:
    hmc_state: HamiltonianMonteCarloState
    hmc_extra: HamiltonianMonteCarloExtra
Hamiltonian Monte Carlo TransitionOperator .