experimental/neutra/neutra_kernel.py
def bootstrap_results(self, state):
    """Trains the bijector and creates initial `previous_kernel_results`.

    The supplied `state` is only used to determine the number of chains to run
    in parallel_iterations

    Args:
      state: `Tensor` or Python `list` of `Tensor`s representing the initial
        state(s) of the Markov chain(s). The first `r` dimensions index
        independent chains, `r = tf.rank(target_log_prob_fn(*state))`.

    Returns:
      kernel_results: Instance of
        `UncalibratedHamiltonianMonteCarloKernelResults` inside
        `MetropolisHastingsResults` inside `TransformedTransitionKernelResults`
        inside `SimpleStepSizeAdaptationResults`.
    """

    def loss():
      q = self._flattened_variational_distribution()
      # TODO(siege): How to seed this?
      samples = q.sample(self.train_batch_size)
      return tf.reduce_mean(
          input_tensor=q.log_prob(samples) -
          self._flattened_target_log_prob(samples),
          axis=-1)

    lr = tf.convert_to_tensor(value=self.learning_rate, dtype=self._dtype)
    dtype = lr.dtype

    learning_rate = tf.compat.v2.optimizers.schedules.PiecewiseConstantDecay(
        list(self.num_train_steps *
             np.array([0.2, 0.8]).astype(dtype.as_numpy_dtype())),
        [lr, lr * 0.1, lr * 0.01])

    opt = tf.compat.v2.optimizers.Adam(learning_rate)

    @tf.function(autograph=False)
    def train_step():
      with tf.GradientTape() as tape:
        loss_val = loss()
      vals = tape.watched_variables()
      grads = tape.gradient(loss_val, vals)
      grads_and_vals = list(zip(grads, vals))
      opt.apply_gradients(grads_and_vals)
      return loss_val

    for step in range(self.num_train_steps):
      loss_val = train_step()
      tf.debugging.assert_all_finite(
          loss_val, 'NeuTra loss is NaN at step {}'.format(step))
      if self.train_debug_fn:
        # pylint: disable=not-callable
        self.train_debug_fn(self, step, loss_val)

    state_parts = tf.nest.flatten(state)
    flat_state_shapes = tf.nest.flatten(self.state_shape)
    batch_shape = tf.shape(input=state_parts[0])[:-flat_state_shapes[0].ndims]

    return self._kernel.bootstrap_results(
        self._flattened_variational_distribution().sample(
            batch_shape, seed=self.seed))
def bootstrap_results(self, state):
    """Trains the bijector and creates initial `previous_kernel_results`.

    The supplied `state` is only used to determine the number of chains to run
    in parallel_iterations

    Args:
      state: `Tensor` or Python `list` of `Tensor`s representing the initial
        state(s) of the Markov chain(s). The first `r` dimensions index
        independent chains, `r = tf.rank(target_log_prob_fn(*state))`.

    Returns:
      kernel_results: Instance of
        `UncalibratedHamiltonianMonteCarloKernelResults` inside
        `MetropolisHastingsResults` inside `TransformedTransitionKernelResults`
        inside `SimpleStepSizeAdaptationResults`.
    """

    def loss():
      q = self._flattened_variational_distribution()
      # TODO(siege): How to seed this?
      samples = q.sample(self.train_batch_size)
      return tf.reduce_mean(
          input_tensor=q.log_prob(samples) -
          self._flattened_target_log_prob(samples),
          axis=-1)

    lr = tf.convert_to_tensor(value=self.learning_rate, dtype=self._dtype)
    dtype = lr.dtype

    learning_rate = tf.compat.v2.optimizers.schedules.PiecewiseConstantDecay(
        list(self.num_train_steps *
             np.array([0.2, 0.8]).astype(dtype.as_numpy_dtype())),
        [lr, lr * 0.1, lr * 0.01])

    opt = tf.compat.v2.optimizers.Adam(learning_rate)

    @tf.function(autograph=False)
    def train_step():
      with tf.GradientTape() as tape:
        loss_val = loss()
      vals = tape.watched_variables()
      grads = tape.gradient(loss_val, vals)
      grads_and_vals = list(zip(grads, vals))
      opt.apply_gradients(grads_and_vals)
      return loss_val

    for step in range(self.num_train_steps):
      loss_val = train_step()
      tf.debugging.assert_all_finite(
          loss_val, 'NeuTra loss is NaN at step {}'.format(step))
      if self.train_debug_fn:
        # pylint: disable=not-callable
        self.train_debug_fn(self, step, loss_val)

    state_parts = tf.nest.flatten(state)
    flat_state_shapes = tf.nest.flatten(self.state_shape)
    batch_shape = tf.shape(input=state_parts[0])[:-flat_state_shapes[0].ndims]

    return self._kernel.bootstrap_results(
        self._flattened_variational_distribution().sample(
            batch_shape, seed=self.seed))
Trains the bijector and creates initial `previous_kernel_results`.

    The supplied `state` is only used to determine the number of chains to run
    in parallel_iterations

    Args:
      state: `Tensor` or Python `list` of `Tensor`s representing the initial
        state(s) of the Markov chain(s). The first `r` dimensions index
        independent chains, `r = tf.rank(target_log_prob_fn(*state))`.

    Returns:
      kernel_results: Instance of
        `UncalibratedHamiltonianMonteCarloKernelResults` inside
        `MetropolisHastingsResults` inside `TransformedTransitionKernelResults`
        inside `SimpleStepSizeAdaptationResults`.