//  [python]
def categorical_bottleneck(dist,
                           vector_quantizer,
                           num_iaf_flows=0,
                           average_categorical_samples=True,
                           use_transformer_for_iaf_parameters=False,
                           sum_over_latents=True,
                           num_samples=1,
                           summary=True):
  """Implements soft EM bottleneck using averaged categorical samples.

  Args:
    dist: Distances between encoder outputs and codebook entries. Negative
      distances are used as categorical logits. A float Tensor of shape
      [batch_size, latent_size, code_size].
    vector_quantizer: An instance of the VectorQuantizer class.
    num_iaf_flows: Number of inverse autoregressive flows.
    average_categorical_samples: Whether to take the average of `num_samples'
      categorical samples as in Roy et al. or to use `num_samples` categorical
      samples to approximate gradient.
    use_transformer_for_iaf_parameters: Whether to use a Transformer instead of
      a lower-triangular mat-mul for generating IAF parameters.
    sum_over_latents: Whether to sum over latent dimension when computing
      entropy.
    num_samples: Number of categorical samples.
    summary: Whether to log entropy histogram.

  Returns:
    one_hot_assignments: Simplex-valued assignments sampled from categorical.
    neg_q_entropy: Negative entropy of categorical distribution.
  """
  latent_size = dist.shape[1]
  x_means_idx = tf.random.categorical(
      logits=tf.reshape(-dist, [-1, FLAGS.num_codes]),
      num_samples=num_samples)
  one_hot_assignments = tf.one_hot(x_means_idx,
                                   depth=vector_quantizer.num_codes)
  one_hot_assignments = tf.reshape(
      one_hot_assignments,
      [-1, latent_size, num_samples, vector_quantizer.num_codes])
  if average_categorical_samples:
    summed_assignments = tf.reduce_sum(one_hot_assignments, axis=2)
    averaged_assignments = tf.reduce_mean(one_hot_assignments, axis=2)
    entropy_dist = tfd.Multinomial(
        total_count=tf.cast(num_samples, tf.float32),
        logits=-dist)
    neg_q_entropy = entropy_dist.log_prob(summed_assignments)
    one_hot_assignments = averaged_assignments[tf.newaxis, Ellipsis]
  else:
    one_hot_assignments = tf.transpose(one_hot_assignments, [2, 0, 1, 3])
    entropy_dist = tfd.OneHotCategorical(logits=-dist, dtype=tf.float32)
    neg_q_entropy = entropy_dist.log_prob(one_hot_assignments)

  if summary:
    tf.summary.histogram("neg_q_entropy_0",
                         tf.reshape(tf.reduce_sum(neg_q_entropy, axis=-1),
                                    [-1]))
  # Perform straight-through.
  class_probs = tf.nn.softmax(-dist[tf.newaxis, Ellipsis])
  one_hot_assignments = class_probs + tf.stop_gradient(
      one_hot_assignments - class_probs)

  # Perform IAF flows
  for flow_num in range(num_iaf_flows):
    with tf.variable_scope("iaf_variables", reuse=tf.AUTO_REUSE):
      # Pad the one_hot_assignments by zeroing out the first latent dimension
      # and shifting the rest down by one (and removing the last dimension).
      shifted_codes = shift_assignments(one_hot_assignments)
      if use_transformer_for_iaf_parameters:
        unconstrained_scale = iaf_scale_from_transformer(
            shifted_codes,
            vector_quantizer.code_size,
            name=str(flow_num))
      else:
        unconstrained_scale = iaf_scale_from_matmul(shifted_codes,
                                                    name=str(flow_num))
      # Initialize scale bias to be log(e/2 - 1) so initial scale + scale_bias
      # is 1.
      initial_scale_bias = tf.fill([latent_size, vector_quantizer.num_codes],
                                   INITIAL_SCALE_BIAS)
      scale_bias = tf.get_variable("scale_bias_" + str(flow_num),
                                   initializer=initial_scale_bias)

    # Since categorical is discrete, we don't need to add inverse log
    # determinant.
    one_hot_assignments, _ = iaf_flow(one_hot_assignments,
                                      unconstrained_scale,
                                      scale_bias,
                                      summary=summary)

  if sum_over_latents:
    neg_q_entropy = tf.reduce_sum(neg_q_entropy, axis=-1)
  neg_q_entropy = tf.reduce_mean(neg_q_entropy)

  return one_hot_assignments, neg_q_entropy