text: # Adapted from https://github.com/tensorflow/models/tree/master/research/fivo
  # Chris J. Maddison*, Dieterich Lawson*, George Tucker*, Nicolas Heess, Mohammad Norouzi, Andriy Mnih, Arnaud Doucet, and Yee Whye Teh. 
  # "Filtering Variational Objectives." NIPS 2017.
  
  """Implementation of sequential Monte Carlo algorithms."""
  
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  import tensorflow as tf
  
  from fivo import utils
  
  
  def ess_criterion(log_weights, unused_t):
    """A criterion that resamples based on effective sample size."""
    num_particles = tf.shape(log_weights)[0]
  
    # Calculate the effective sample size.
    ess_num = 2 * tf.reduce_logsumexp(log_weights, axis=0)
    ess_denom = tf.reduce_logsumexp(2 * log_weights, axis=0)
    log_ess = ess_num - ess_denom
  
    return log_ess <= tf.log(tf.to_float(num_particles) / 2.0)
  
  
  def never_resample_criterion(log_weights, unused_t):
    """A criterion that never resamples."""
    batch_size = tf.shape(log_weights)[1]
  
    return tf.cast(tf.zeros([batch_size]), tf.bool)
  
  
  def always_resample_criterion(log_weights, unused_t):
    """A criterion resamples at every timestep."""
    batch_size = tf.shape(log_weights)[1]
  
    return tf.cast(tf.ones([batch_size]), tf.bool)
  
  
  def multinomial_resampling(log_weights, states, num_particles, batch_size,
                             random_seed=None):
    """Resample states with multinomial resampling.
  
    Args:
      log_weights: A [num_particles, batch_size] Tensor representing a batch
        of batch_size logits for num_particles-ary Categorical distribution.
      states: A nested list of [batch_size*num_particles, data_size] Tensors that
        will be resampled from the groups of every num_particles-th row.
      num_particles: The number of particles/samples.
      batch_size: The batch size.
      random_seed: The random seed to pass to the resampling operations in
        the particle filter. Mainly useful for testing.
  
    Returns:
      resampled_states: A nested list of [batch_size*num_particles, data_size]
        Tensors resampled via multinomial sampling.
    """
  
    # Calculate the ancestor indices via resampling. Because we maintain the
    # log unnormalized weights, we pass the weights in as logits, allowing
    # the distribution object to apply a softmax and normalize them.
    resampling_parameters = tf.transpose(log_weights, perm=[1, 0])
    resampling_dist = tf.contrib.distributions.Categorical(
        logits=resampling_parameters)
    ancestors = tf.stop_gradient(
        resampling_dist.sample(sample_shape=num_particles, seed=random_seed))
  
    # Because the batch is flattened, we must modify ancestor_inds to index the
    # proper samples. The particles in the ith filter are distributed every
    # batch_size rows in the batch, and offset i rows from the top. So, to
    # correct the indices we multiply by the batch_size and add the proper offset.
    # Crucially, when ancestor_inds is flattened the layout of the batch is
    # maintained.
    offset = tf.expand_dims(tf.range(batch_size), 0)
    ancestor_inds = tf.reshape(ancestors * batch_size + offset, [-1])
  
    resampled_states = utils.gather_tensors(states, ancestor_inds)
  
    return resampled_states
  
  
  def _blend_tensor(blending_weights, tensor, num_particles, batch_size):
    """Blend tensor according to the weights.
  
    The first dimension of tensor is actually a 2d index compacted to a 1d
    index and similarly for blended_tensor. So if we index these Tensors
    by [(i, j), k], then
  
      blended_tensor[(i, j), k] =
        sum_l tensor[(l, j), :] * blending_weights[i, j, l].
  
    Args:
      blending_weights: [num_particles, batch_size, num_particles] weights where
        the indices represent [sample index, batch index, blending weight index].
      tensor: [num_particles * batch_size, state_dim] Tensor to be blended.
      num_particles: The number of particles/samples.
      batch_size: The batch size.
  
    Returns:
      blended_tensor: [num_particles*batch_size, state_dim] blended Tensor.
    """
    # tensor is currently [num_particles * batch_size, state_dim], so we reshape
    # it to [num_particles, batch_size, state_dim]. Then, transpose it to
    # [batch_size, state_size, num_particles].
    tensor = tf.transpose(
        tf.reshape(tensor, [num_particles, batch_size, -1]), perm=[1, 2, 0])
  
    blending_weights = tf.transpose(blending_weights, perm=[1, 2, 0])
    # blendeding_weights is [batch index, blending weight index, sample index].
    # Multiplying these gives a matrix of size [batch_size, state_size,
    # num_particles].
    tensor = tf.matmul(tensor, blending_weights)
    # transpose the tensor to be [num_particles, batch_size, state_size]
    # and then reshape it to match the original format.
    tensor = tf.reshape(tf.transpose(tensor, perm=[2, 0, 1]),
                        [num_particles*batch_size, -1])
  
    return tensor
  
  
  def relaxed_resampling(log_weights, states, num_particles, batch_size,
                         temperature=0.5, random_seed=None):
    """Resample states with relaxed resampling.
  
    Draw soft "ancestors" using the Gumbel-Softmax distribution.
  
    Args:
      log_weights: A [num_particles, batch_size] Tensor representing a batch
        of batch_size logits for num_particles-ary Categorical distribution.
      states: A nested list of [batch_size * num_particles, d] Tensors that will
        be resampled from the groups of every num_particles-th row.
      num_particles: The number of particles/samples.
      batch_size: The batch size.
      temperature: The temperature used for the relaxed one hot distribution.
      random_seed: The random seed to pass to the resampling operations in
        the particle filter. Mainly useful for testing.
  
    Returns:
      resampled_states: A nested list of [batch_size * num_particles, d]
        Tensors resampled via multinomial sampling.
    """
    # log_weights are [num_particles, batch_size], so we transpose to get a
    # set of batch_size distributions over [0, num_particles).
    resampling_parameters = tf.transpose(log_weights, perm=[1, 0])
    resampling_dist = tf.contrib.distributions.RelaxedOneHotCategorical(
        temperature,
        logits=resampling_parameters)
  
    # Sample num_particles samples from the distribution, resulting in a
    # [num_particles, batch_size, num_particles] Tensor that represents a set of
    # [num_particles, batch_size] blending weights. The dimensions represent
    # [particle index, batch index, blending weight index].
    ancestors = resampling_dist.sample(sample_shape=num_particles,
                                       seed=random_seed)
    def map_fn(tensor):
      return _blend_tensor(ancestors, tensor, num_particles, batch_size)
  
    resampled_states = utils.map_nested(map_fn, states)
  
    return resampled_states
  
  
  def smc(
      transition_fn,
      num_steps,
      num_particles=1,
      resampling_criterion=ess_criterion,
      resampling_fn=multinomial_resampling,
      loop_fn=None,
      parallel_iterations=30,
      swap_memory=True):
    """Run a sequential Monte Carlo (SMC) algorithm.
  
    This method runs an SMC algorithm that evolves systems of particles
    using the supplied transition function for the specified number of steps. The
    particles are optionally resampled using resampling_fn when indicated by
    resampling_criterion.
  
    Args:
      transition_fn: A callable that propogates a batch of particles one step.
        Must accept as arguments a batch of particle states and the current
        timestep. Must return the particle states one timestep in the future, the
        incremental weights of each particle as a [num_samples*batch_size] float
        Tensor, and optionally a set of arguments to pass to the loop_fn. If
        the loop args are not provided, they will be set to None. Before the
        first timestep transition_fn will be called with the arguments None, -1
        and should return the initial particle states.
      num_steps: A [batch_size] Tensor of ints representing the number of steps
        to run each filter for.
      num_particles: A scalar int, the number of particles to use in each filter.
      resampling_criterion: The resampling criterion to use for this particle
        filter. Must accept the current log weights and timestep and
        return a boolean Tensor of shape [batch_size] indicating whether each
        particle filter should resample. See ess_criterion and related functions
        for examples. When resampling_criterion is never_resample_criterion,
        resampling_fn is ignored and never called.
      resampling_fn: A callable that performs the resampling operation. Must
        accept as arguments the log weights, particle states, num_particles,
        and batch_size and return the resampled particle states. See
        multinomial_resampling and relaxed_resampling for examples.
      loop_fn: A callable that performs operations on the weights and
        particle states, useful for accumulating and processing state that
        shouldn't be resampled. At each timestep after (possibly) resampling
        loop_fn will be called with the previous loop_state, a set of arguments
        produced by transition_fn called loop_args, the resampled particle states,
        the current log weights as [num_particles, batch_size] float Tensor, a
        [batch_size] float Tensor representing whether or not each filter
        resampled, the current mask indicating which filters are active, and the
        current timestep. It must return the next loop state. Before the first
        timestep loop_fn will be called with the arguments None, None, None, None,
        -1 and must return the initial loop state. The loop state can be a
        possibly nested structure of Tensors and TensorArrays.
      parallel_iterations: The number of parallel iterations to use for the
        internal while loop. Note that values greater than 1 can introduce
        non-determinism even when resampling is deterministic.
      swap_memory: Whether GPU-CPU memory swapping should be enabled for the
        internal while loop.
  
    Returns:
      log_z_hat: A Tensor of shape [batch_size] containing an estimate of the log
        normalizing constant that converts between the unormalized target
        distribution (as defined by the weights) and the true target distribution.
      log_weights: A Tensor of shape [max_num_steps, batch_size, num_particles]
        containing the log weights at each timestep of the particle filter.
        Will not be valid for timesteps past the supplied num_steps.
      resampled: A float Tensor of shape [max_num_steps, batch_size] indicating
        when the particle filters resampled. Will be 1.0 on timesteps when
        resampling occurred and 0.0 on timesteps when it did not.
      final_loop_state: The final state returned by loop_fn. If loop_fn is None
        then 0 will be returned.
    """
    # batch_size represents the number of particle filters running in parallel.
    batch_size = tf.shape(num_steps)[0]
    # Create a TensorArray where element t is the [num_particles*batch_size]
    # sequence mask for timestep t.
    max_num_steps = tf.reduce_max(num_steps)
    seq_mask = tf.transpose(
        tf.sequence_mask(num_steps, maxlen=max_num_steps, dtype=tf.float32),
        perm=[1, 0])
  
    seq_mask = tf.tile(seq_mask, [1, num_particles])
    mask_ta = tf.TensorArray(seq_mask.dtype,
                             max_num_steps,
                             name='mask_ta')
    mask_ta = mask_ta.unstack(seq_mask)
  
    # Initialize the state.
    t0 = tf.constant(0, tf.int32)
    init_particle_state = transition_fn(None, -1)
  
    def transition(*args):
      transition_outs = transition_fn(*args)
      if len(transition_outs) == 2:
        return transition_outs + (None,)
      else:
        return transition_outs
  
    if loop_fn is None:
      loop_fn = lambda *args: 0
  
    init_loop_state = loop_fn(None, None, None, None, None, None, -1)
    init_states = (init_particle_state, init_loop_state)
    ta_names = ['log_weights', 'resampled']
    tas = [tf.TensorArray(tf.float32, max_num_steps, name='%s_ta' % n)
           for n in ta_names]
    log_weights_acc = tf.zeros([num_particles, batch_size], dtype=tf.float32)
    log_z_hat_acc = tf.zeros([batch_size], dtype=tf.float32)
  
    def while_predicate(t, *unused_args):
      return t < max_num_steps
  
    def while_step(t, state, tas, log_weights_acc, log_z_hat_acc):
      """Implements one timestep of the particle filter."""
      particle_state, loop_state = state
      cur_mask = utils.read_tas(mask_ta, t)
      # Propagate the particles one step.
      log_alpha, new_particle_state, loop_args = transition(particle_state, t)
      # Update the current weights with the incremental weights.
      log_alpha *= cur_mask
      log_alpha = tf.reshape(log_alpha, [num_particles, batch_size])
      log_weights_acc += log_alpha
  
      should_resample = resampling_criterion(log_weights_acc, t)
  
      if resampling_criterion == never_resample_criterion:
        resampled = tf.to_float(should_resample)
      else:
        # Compute the states as if we did resample.
        resampled_states = resampling_fn(
            log_weights_acc,
            new_particle_state,
            num_particles,
            batch_size)
        # Decide whether or not we should resample; don't resample if we are past
        # the end of a sequence.
        should_resample = tf.logical_and(should_resample,
                                         cur_mask[:batch_size] > 0.)
        float_should_resample = tf.to_float(should_resample)
        new_particle_state = utils.where_tensors(
            tf.tile(should_resample, [num_particles]),
            resampled_states,
            new_particle_state)
        resampled = float_should_resample
  
      new_loop_state = loop_fn(loop_state, loop_args, new_particle_state,
                               log_weights_acc, resampled, cur_mask, t)
      # Update log Z hat.
      log_z_hat_update = tf.reduce_logsumexp(
          log_weights_acc, axis=0) - tf.log(tf.to_float(num_particles))
      # If it is the last timestep, always add the update.
      log_z_hat_acc += tf.cond(t < max_num_steps - 1,
                               lambda: log_z_hat_update * resampled,
                               lambda: log_z_hat_update)
      # Update the TensorArrays before we reset the weights so that we capture
      # the incremental weights and not zeros.
      ta_updates = [log_weights_acc, resampled]
      new_tas = [ta.write(t, x) for ta, x in zip(tas, ta_updates)]
      # For the particle filters that resampled, reset weights to zero.
      log_weights_acc *= (1. - tf.tile(resampled[tf.newaxis, :],
                                       [num_particles, 1]))
      new_state = (new_particle_state, new_loop_state)
      
      return t + 1, new_state, new_tas, log_weights_acc, log_z_hat_acc
  
    _, final_state, tas, _, log_z_hat = tf.while_loop(
        while_predicate,
        while_step,
        loop_vars=(t0, init_states, tas, log_weights_acc, log_z_hat_acc),
        parallel_iterations=parallel_iterations,
        swap_memory=swap_memory)
  
    log_weights, resampled = [x.stack() for x in tas]
    log_weights = tf.transpose(log_weights, perm=[0, 2, 1])
    final_particle_state, final_loop_state = final_state
  
    return (log_z_hat, log_weights, resampled,
            final_particle_state, final_loop_state)
  