jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def preprocess_for_train(image,
                         output_height,
                         output_width,
                         resize_side_min=_RESIZE_SIDE_MIN,
                         resize_side_max=_RESIZE_SIDE_MAX,
                         use_grayscale=False):
  """Preprocesses the given image for training.

  Note that the actual resizing scale is sampled from
    [`resize_size_min`, `resize_size_max`].

  Args:
    image: A `Tensor` representing an image of arbitrary size.
    output_height: The height of the image after preprocessing.
    output_width: The width of the image after preprocessing.
    resize_side_min: The lower bound for the smallest side of the image for
      aspect-preserving resizing.
    resize_side_max: The upper bound for the smallest side of the image for
      aspect-preserving resizing.
    use_grayscale: Whether to convert the image from RGB to grayscale.

  Returns:
    A preprocessed image.
  """
  resize_side = tf.random_uniform(
      [], minval=resize_side_min, maxval=resize_side_max+1, dtype=tf.int32)

  image = _aspect_preserving_resize(image, resize_side)
  image = _random_crop([image], output_height, output_width)[0]
  image.set_shape([output_height, output_width, 3])
  image = tf.cast(image, tf.float32)
  if use_grayscale:
    image = tf.image.rgb_to_grayscale(image)
  image = tf.image.random_flip_left_right(image)
  return _mean_image_subtraction(image, [_R_MEAN, _G_MEAN, _B_MEAN])
//  [python]
def generate_optimal_weights(dimension, distance):
  """Makes random optimal weights for use in make_features."""
  direction = np.random.normal(0.0, 1.0, [dimension])
  direction = direction / np.linalg.norm(direction)
  weights = distance * direction
  weights = np.reshape(weights, [dimension, 1])
  return tf.to_float(tf.constant(weights))
//  [python]
def generate_optimal_weights(dimension, distance):
  """Makes random optimal weights for use in make_features."""
  direction = np.random.normal(0.0, 1.0, [dimension])
  direction = direction / np.linalg.norm(direction)
  weights = distance * direction
  weights = np.reshape(weights, [dimension, 1])
  return tf.cast(tf.constant(weights), tf.float32)
//  [python]
def input_fn_builder(tfrecord_file,
                     seq_length,
                     is_training,
                     batch_size,
                     drop_remainder=True):
  """Creates an `input_fn` closure to be passed to Estimator."""
  logging.info(
      "Creating input fun with batch_size: {} and drop remainder: {}".format(
          str(batch_size), str(drop_remainder)))
  name_to_features = {
      "input_ids": tf.FixedLenFeature([seq_length], tf.int64),
      "input_mask": tf.FixedLenFeature([seq_length], tf.int64),
      "segment_ids": tf.FixedLenFeature([seq_length], tf.int64),
      "score": tf.FixedLenFeature([], tf.float32)
  }

  def _decode_record(record, name_to_features):
    """Decodes a record to a TensorFlow example."""
    example = tf.parse_single_example(record, name_to_features)
    # tf.Example only supports tf.int64, but the TPU only supports tf.int32.
    # So cast all int64 to int32.
    for name in list(example.keys()):
      t = example[name]
      if t.dtype == tf.int64:
        t = tf.to_int32(t)
      example[name] = t
    return example

  def input_fn(params):  # pylint: disable=unused-argument
    """Acutal data generator."""
    d = tf.data.TFRecordDataset(tfrecord_file)
    if is_training:
      d = d.repeat()
      d = d.shuffle(buffer_size=FLAGS.shuffle_buffer_size)
    d = d.map(lambda record: _decode_record(record, name_to_features))
    d = d.batch(batch_size=batch_size, drop_remainder=drop_remainder)
    return d

  return input_fn
//  [python]
def input_fn_builder(tfrecord_file,
                     seq_length,
                     is_training,
                     batch_size,
                     drop_remainder=True):
  """Creates an `input_fn` closure to be passed to Estimator."""
  logging.info(
      "Creating input fun with batch_size: {} and drop remainder: {}".format(
          str(batch_size), str(drop_remainder)))
  name_to_features = {
      "input_ids": tf.FixedLenFeature([seq_length], tf.int64),
      "input_mask": tf.FixedLenFeature([seq_length], tf.int64),
      "segment_ids": tf.FixedLenFeature([seq_length], tf.int64),
      "score": tf.FixedLenFeature([], tf.float32)
  }

  def _decode_record(record, name_to_features):
    """Decodes a record to a TensorFlow example."""
    example = tf.parse_single_example(record, name_to_features)
    # tf.Example only supports tf.int64, but the TPU only supports tf.int32.
    # So cast all int64 to int32.
    for name in list(example.keys()):
      t = example[name]
      if t.dtype == tf.int64:
        t = tf.cast(t, tf.int32)
      example[name] = t
    return example

  def input_fn(params):  # pylint: disable=unused-argument
    """Acutal data generator."""
    d = tf.data.TFRecordDataset(tfrecord_file)
    if is_training:
      d = d.repeat()
      d = d.shuffle(buffer_size=FLAGS.shuffle_buffer_size)
    d = d.map(lambda record: _decode_record(record, name_to_features))
    d = d.batch(batch_size=batch_size, drop_remainder=drop_remainder)
    return d

  return input_fn
//  [python]
def gaussian_blur(image, kernel_size, sigma, padding='SAME'):
  """Blurs the given image with separable convolution.


  Args:
    image: Tensor of shape [height, width, channels] and dtype float to blur.
    kernel_size: Integer Tensor for the size of the blur kernel. This is should
      be an odd number. If it is an even number, the actual kernel size will be
      size + 1.
    sigma: Sigma value for gaussian operator.
    padding: Padding to use for the convolution. Typically 'SAME' or 'VALID'.

  Returns:
    A Tensor representing the blurred image.
  """
  radius = tf.to_int32(kernel_size / 2)
  kernel_size = radius * 2 + 1
  x = tf.to_float(tf.range(-radius, radius + 1))
  blur_filter = tf.exp(
      -tf.pow(x, 2.0) / (2.0 * tf.pow(tf.to_float(sigma), 2.0)))
  blur_filter /= tf.reduce_sum(blur_filter)
  # One vertical and one horizontal filter.
  blur_v = tf.reshape(blur_filter, [kernel_size, 1, 1, 1])
  blur_h = tf.reshape(blur_filter, [1, kernel_size, 1, 1])
  num_channels = tf.shape(image)[-1]
  blur_h = tf.tile(blur_h, [1, 1, num_channels, 1])
  blur_v = tf.tile(blur_v, [1, 1, num_channels, 1])
  expand_batch_dim = image.shape.ndims == 3
  if expand_batch_dim:
    # Tensorflow requires batched input to convolutions, which we can fake with
    # an extra dimension.
    image = tf.expand_dims(image, axis=0)
  blurred = tf.nn.depthwise_conv2d(
      image, blur_h, strides=[1, 1, 1, 1], padding=padding)
  blurred = tf.nn.depthwise_conv2d(
      blurred, blur_v, strides=[1, 1, 1, 1], padding=padding)
  if expand_batch_dim:
    blurred = tf.squeeze(blurred, axis=0)
  return blurred
//  [python]
def gaussian_blur(image, kernel_size, sigma, padding='SAME'):
  """Blurs the given image with separable convolution.


  Args:
    image: Tensor of shape [height, width, channels] and dtype float to blur.
    kernel_size: Integer Tensor for the size of the blur kernel. This is should
      be an odd number. If it is an even number, the actual kernel size will be
      size + 1.
    sigma: Sigma value for gaussian operator.
    padding: Padding to use for the convolution. Typically 'SAME' or 'VALID'.

  Returns:
    A Tensor representing the blurred image.
  """
  radius = tf.cast(kernel_size / 2, tf.int32)
  kernel_size = radius * 2 + 1
  x = tf.cast(tf.range(-radius, radius + 1), tf.float32)
  blur_filter = tf.exp(
      -tf.pow(x, 2.0) / (2.0 * tf.pow(tf.cast(sigma, tf.float32), 2.0)))
  blur_filter /= tf.reduce_sum(blur_filter)
  # One vertical and one horizontal filter.
  blur_v = tf.reshape(blur_filter, [kernel_size, 1, 1, 1])
  blur_h = tf.reshape(blur_filter, [1, kernel_size, 1, 1])
  num_channels = tf.shape(image)[-1]
  blur_h = tf.tile(blur_h, [1, 1, num_channels, 1])
  blur_v = tf.tile(blur_v, [1, 1, num_channels, 1])
  expand_batch_dim = image.shape.ndims == 3
  if expand_batch_dim:
    # Tensorflow requires batched input to convolutions, which we can fake with
    # an extra dimension.
    image = tf.expand_dims(image, axis=0)
  blurred = tf.nn.depthwise_conv2d(
      image, blur_h, strides=[1, 1, 1, 1], padding=padding)
  blurred = tf.nn.depthwise_conv2d(
      blurred, blur_v, strides=[1, 1, 1, 1], padding=padding)
  if expand_batch_dim:
    blurred = tf.squeeze(blurred, axis=0)
  return blurred
//  [python]
def DenseMNISTNoise(x=None, z=None, decoder=None, return_means=True):
  """Models P(x | decoder(z))"""
  decoding = decoder(z)
  bernoulli = IndependentBernouli3D(decoding)

  if x is None:
    if return_means:
      x = bernoulli.mean()
    else:
      x = tf.to_float(bernoulli.sample())

  return x, bernoulli.log_prob(x)
//  [python]
def DenseMNISTNoise(x=None, z=None, decoder=None, return_means=True):
  """Models P(x | decoder(z))"""
  decoding = decoder(z)
  bernoulli = IndependentBernouli3D(decoding)

  if x is None:
    if return_means:
      x = bernoulli.mean()
    else:
      x = tf.cast(bernoulli.sample(), tf.float32)

  return x, bernoulli.log_prob(x)
//  [python]
def get_masked_mention_word_scores(self, encoded_doc, span_starts, span_ends):
      num_words = util.shape(encoded_doc, 0) # T
      num_c = util.shape(span_starts, 0) # NC
      doc_range = tf.tile(tf.expand_dims(tf.range(0, num_words), 0), [num_c, 1]) # [K, T]
      mention_mask = tf.logical_and(doc_range >= tf.expand_dims(span_starts, 1), doc_range <= tf.expand_dims(span_ends, 1)) #[K, T]
      with tf.variable_scope("mention_word_attn"):
        word_attn = tf.squeeze(util.projection(encoded_doc, 1, initializer=tf.truncated_normal_initializer(stddev=0.02)), 1)
      mention_word_attn = tf.nn.softmax(tf.log(tf.to_float(mention_mask)) + tf.expand_dims(word_attn, 0))
      return mention_word_attn
//  [python]
def get_masked_mention_word_scores(self, encoded_doc, span_starts, span_ends):
      num_words = util.shape(encoded_doc, 0) # T
      num_c = util.shape(span_starts, 0) # NC
      doc_range = tf.tile(tf.expand_dims(tf.range(0, num_words), 0), [num_c, 1]) # [K, T]
      mention_mask = tf.logical_and(doc_range >= tf.expand_dims(span_starts, 1), doc_range <= tf.expand_dims(span_ends, 1)) #[K, T]
      with tf.variable_scope("mention_word_attn"):
        word_attn = tf.squeeze(util.projection(encoded_doc, 1, initializer=tf.truncated_normal_initializer(stddev=0.02)), 1)
      mention_word_attn = tf.nn.softmax(tf.log(tf.cast(mention_mask, tf.float32)) + tf.expand_dims(word_attn, 0))
      return mention_word_attn
//  [python]
def AdjustedStepSize(self):
    if self._step_size_warmup > 0:
      global_step = tf.train.get_or_create_global_step()
      max_step = self._init_step_size * tf.to_float(
          global_step) / self._step_size_warmup
      return tf.where(global_step > self._step_size_warmup, self._step_size,
                      tf.minimum(max_step, self._step_size))
    else:
      return self._step_size
//  [python]
def AdjustedStepSize(self):
    if self._step_size_warmup > 0:
      global_step = tf.train.get_or_create_global_step()
      max_step = self._init_step_size * tf.cast(
          global_step, tf.float32) / self._step_size_warmup
      return tf.where(global_step > self._step_size_warmup, self._step_size,
                      tf.minimum(max_step, self._step_size))
    else:
      return self._step_size
//  [python]
def DenseCIFAR10TNoise(x=None, z=None, decoder=None, return_means=True, uniform_scale=False, logistic_impl="mine"):
  """Models P(x | decoder(z))"""
  decoding = decoder(z)
  if uniform_scale:
    scale = tf.get_variable("scale", initializer=1.0)
    scales = tf.reshape(scale, [1, 1, 1])
  else:
    scales = tf.get_variable(
        "scales", [32, 32, 3], initializer=tf.ones_initializer())

  if logistic_impl == "mine":
    disc_logistic = IndependentDiscreteLogistic3D(decoding, tf.nn.softplus(scales))
  elif logistic_impl == "kingma":
    disc_logistic = IndependentDiscreteLogistic3D2(decoding, tf.nn.softplus(scales))

  if x is None:
    x = tf.to_float(disc_logistic.sample())

  return x, disc_logistic.log_prob(x)
//  [python]
def DenseCIFAR10TNoise(x=None, z=None, decoder=None, return_means=True, uniform_scale=False, logistic_impl="mine"):
  """Models P(x | decoder(z))"""
  decoding = decoder(z)
  if uniform_scale:
    scale = tf.get_variable("scale", initializer=1.0)
    scales = tf.reshape(scale, [1, 1, 1])
  else:
    scales = tf.get_variable(
        "scales", [32, 32, 3], initializer=tf.ones_initializer())

  if logistic_impl == "mine":
    disc_logistic = IndependentDiscreteLogistic3D(decoding, tf.nn.softplus(scales))
  elif logistic_impl == "kingma":
    disc_logistic = IndependentDiscreteLogistic3D2(decoding, tf.nn.softplus(scales))

  if x is None:
    x = tf.cast(disc_logistic.sample(), tf.float32)

  return x, disc_logistic.log_prob(x)
//  [python]
def input_fn_builder(input_file, seq_length, is_training, drop_remainder):
  """Creates an `input_fn` closure to be passed to TPUEstimator."""

  name_to_features = {
      "unique_ids": tf.FixedLenFeature([], tf.int64),
      "input_ids": tf.FixedLenFeature([seq_length], tf.int64),
      "input_mask": tf.FixedLenFeature([seq_length], tf.int64),
      "segment_ids": tf.FixedLenFeature([seq_length], tf.int64),
  }

  if is_training:
    name_to_features["start_positions"] = tf.FixedLenFeature([], tf.int64)
    name_to_features["end_positions"] = tf.FixedLenFeature([], tf.int64)

  def _decode_record(record, name_to_features):
    """Decodes a record to a TensorFlow example."""
    example = tf.parse_single_example(record, name_to_features)

    # tf.Example only supports tf.int64, but the TPU only supports tf.int32.
    # So cast all int64 to int32.
    for name in list(example.keys()):
      t = example[name]
      if t.dtype == tf.int64:
        t = tf.to_int32(t)
      example[name] = t

    return example

  def input_fn(params):
    """The actual input function."""
    batch_size = params["batch_size"]

    # For training, we want a lot of parallel reading and shuffling.
    # For eval, we want no shuffling and parallel reading doesn't matter.
    d = tf.data.TFRecordDataset(input_file)
    if is_training:
      d = d.repeat()
      d = d.shuffle(buffer_size=100)

    d = d.apply(
        tf.data.experimental.map_and_batch(
            lambda record: _decode_record(record, name_to_features),
            batch_size=batch_size,
            drop_remainder=drop_remainder))

    return d

  return input_fn
//  [python]
def input_fn_builder(input_file, seq_length, is_training, drop_remainder):
  """Creates an `input_fn` closure to be passed to TPUEstimator."""

  name_to_features = {
      "unique_ids": tf.FixedLenFeature([], tf.int64),
      "input_ids": tf.FixedLenFeature([seq_length], tf.int64),
      "input_mask": tf.FixedLenFeature([seq_length], tf.int64),
      "segment_ids": tf.FixedLenFeature([seq_length], tf.int64),
  }

  if is_training:
    name_to_features["start_positions"] = tf.FixedLenFeature([], tf.int64)
    name_to_features["end_positions"] = tf.FixedLenFeature([], tf.int64)

  def _decode_record(record, name_to_features):
    """Decodes a record to a TensorFlow example."""
    example = tf.parse_single_example(record, name_to_features)

    # tf.Example only supports tf.int64, but the TPU only supports tf.int32.
    # So cast all int64 to int32.
    for name in list(example.keys()):
      t = example[name]
      if t.dtype == tf.int64:
        t = tf.cast(t, tf.int32)
      example[name] = t

    return example

  def input_fn(params):
    """The actual input function."""
    batch_size = params["batch_size"]

    # For training, we want a lot of parallel reading and shuffling.
    # For eval, we want no shuffling and parallel reading doesn't matter.
    d = tf.data.TFRecordDataset(input_file)
    if is_training:
      d = d.repeat()
      d = d.shuffle(buffer_size=100)

    d = d.apply(
        tf.data.experimental.map_and_batch(
            lambda record: _decode_record(record, name_to_features),
            batch_size=batch_size,
            drop_remainder=drop_remainder))

    return d

  return input_fn
//  [python]
def LearningRate(train_size, global_step, schedule = "hoffman", warmup_steps=0):
  if schedule == "hoffman":
    base = tf.train.piecewise_constant(
        global_step, [train_size * 500 // TRAIN_BATCH], [1e-3, 1e-4])
  elif schedule == "new":
    base = tf.train.piecewise_constant(
        global_step,
        [train_size * 500 // TRAIN_BATCH, train_size * 800 // TRAIN_BATCH],
        [1e-3, 1e-4, 1e-5])
  elif schedule == "new_gentle":
    base = tf.train.piecewise_constant(
        global_step,
        [train_size * 500 // TRAIN_BATCH, train_size * 800 // TRAIN_BATCH],
        [0.5e-3, 1e-4, 1e-5])
  elif schedule == "fast":
    base = tf.train.piecewise_constant(
        global_step,
        [train_size * 800 // TRAIN_BATCH],
        [1e-2, 1e-5])
  else:
    raise ValueError("Invalid schedule: " + schedule)
  if warmup_steps == 0:
    return base
  else:
    return tf.minimum(base * tf.to_float(global_step) / warmup_steps, base)
//  [python]
def LearningRate(train_size, global_step, schedule = "hoffman", warmup_steps=0):
  if schedule == "hoffman":
    base = tf.train.piecewise_constant(
        global_step, [train_size * 500 // TRAIN_BATCH], [1e-3, 1e-4])
  elif schedule == "new":
    base = tf.train.piecewise_constant(
        global_step,
        [train_size * 500 // TRAIN_BATCH, train_size * 800 // TRAIN_BATCH],
        [1e-3, 1e-4, 1e-5])
  elif schedule == "new_gentle":
    base = tf.train.piecewise_constant(
        global_step,
        [train_size * 500 // TRAIN_BATCH, train_size * 800 // TRAIN_BATCH],
        [0.5e-3, 1e-4, 1e-5])
  elif schedule == "fast":
    base = tf.train.piecewise_constant(
        global_step,
        [train_size * 800 // TRAIN_BATCH],
        [1e-2, 1e-5])
  else:
    raise ValueError("Invalid schedule: " + schedule)
  if warmup_steps == 0:
    return base
  else:
    return tf.minimum(base * tf.cast(global_step, tf.float32) / warmup_steps, base)
//  [python]
def input_fn_builder(input_file_pattern,
                     model_config,
                     long_seq_len,
                     global_seq_len,
                     is_training,
                     drop_remainder,
                     num_cpu_threads=4,
                     candidate_ignore_hard_g2l=True,
                     query_ignore_hard_g2l=True,
                     enable_l2g_linking=True):
  """Creates an `input_fn` closure to be passed to TPUEstimator."""

  def _decode_record(record, name_to_features):
    """Decodes a record to a TensorFlow example."""
    example = tf.parse_single_example(record, name_to_features)

    # tf.Example only supports tf.int64, but the TPU only supports tf.int32.
    # So cast all int64 to int32.
    for name in list(example.keys()):
      t = example[name]
      if t.dtype == tf.int64:
        t = tf.to_int32(t)
      example[name] = t

    return example

  def input_fn(params):
    """The actual input function."""
    name_to_features = get_model_schema(long_seq_len, global_seq_len)
    batch_size = params["batch_size"]

    # For training, we want a lot of parallel reading and shuffling.
    # For eval, we want no shuffling and parallel reading doesn't matter.
    d = tf.data.Dataset.list_files(input_file_pattern, shuffle=is_training)

    # `sloppy` mode means that the interleaving is not exact. This adds
    # even more randomness to the training pipeline.
    d = d.apply(
        tf.data.experimental.parallel_interleave(
            tf.data.TFRecordDataset, cycle_length=num_cpu_threads,
            sloppy=False))

    # For training, we want a lot of parallel reading and shuffling.
    # For eval, we want no shuffling and parallel reading doesn't matter.
    if is_training:
      d = d.shuffle(buffer_size=100).repeat()

    d = d.apply(
        tf.data.experimental.map_and_batch(
            lambda record: _decode_record(record, name_to_features),
            batch_size=batch_size,
            drop_remainder=drop_remainder))
    d = d.prefetch(tf.data.experimental.AUTOTUNE)
    d = d.map(
        functools.partial(
            _add_side_input_features,
            model_config=model_config,
            candidate_ignore_hard_g2l=candidate_ignore_hard_g2l,
            query_ignore_hard_g2l=query_ignore_hard_g2l,
            enable_l2g_linking=enable_l2g_linking),
        tf.data.experimental.AUTOTUNE)

    return d.prefetch(tf.data.experimental.AUTOTUNE)

  return input_fn
//  [python]
def input_fn_builder(input_file_pattern,
                     model_config,
                     long_seq_len,
                     global_seq_len,
                     is_training,
                     drop_remainder,
                     num_cpu_threads=4,
                     candidate_ignore_hard_g2l=True,
                     query_ignore_hard_g2l=True,
                     enable_l2g_linking=True):
  """Creates an `input_fn` closure to be passed to TPUEstimator."""

  def _decode_record(record, name_to_features):
    """Decodes a record to a TensorFlow example."""
    example = tf.parse_single_example(record, name_to_features)

    # tf.Example only supports tf.int64, but the TPU only supports tf.int32.
    # So cast all int64 to int32.
    for name in list(example.keys()):
      t = example[name]
      if t.dtype == tf.int64:
        t = tf.cast(t, tf.int32)
      example[name] = t

    return example

  def input_fn(params):
    """The actual input function."""
    name_to_features = get_model_schema(long_seq_len, global_seq_len)
    batch_size = params["batch_size"]

    # For training, we want a lot of parallel reading and shuffling.
    # For eval, we want no shuffling and parallel reading doesn't matter.
    d = tf.data.Dataset.list_files(input_file_pattern, shuffle=is_training)

    # `sloppy` mode means that the interleaving is not exact. This adds
    # even more randomness to the training pipeline.
    d = d.apply(
        tf.data.experimental.parallel_interleave(
            tf.data.TFRecordDataset, cycle_length=num_cpu_threads,
            sloppy=False))

    # For training, we want a lot of parallel reading and shuffling.
    # For eval, we want no shuffling and parallel reading doesn't matter.
    if is_training:
      d = d.shuffle(buffer_size=100).repeat()

    d = d.apply(
        tf.data.experimental.map_and_batch(
            lambda record: _decode_record(record, name_to_features),
            batch_size=batch_size,
            drop_remainder=drop_remainder))
    d = d.prefetch(tf.data.experimental.AUTOTUNE)
    d = d.map(
        functools.partial(
            _add_side_input_features,
            model_config=model_config,
            candidate_ignore_hard_g2l=candidate_ignore_hard_g2l,
            query_ignore_hard_g2l=query_ignore_hard_g2l,
            enable_l2g_linking=enable_l2g_linking),
        tf.data.experimental.AUTOTUNE)

    return d.prefetch(tf.data.experimental.AUTOTUNE)

  return input_fn
//  [python]
def AIS(self, images, num_chains):

    def ProposalLogProbFn(*z):
      if self._use_q_z_for_ais:
        _, log_p_z, _ = self._recog(images, z=z)
      else:
        _, log_p_z = self._prior(z)
      return tf.add_n(log_p_z)

    def TargetLogProbFn(*z):
      _, log_p_z = self._prior(z)
      _, log_p_x_z = self._noise(images, z)
      return tf.add_n(log_p_z) + log_p_x_z

    images = tf.tile(images, [num_chains, 1, 1, 1])
    if self._use_q_z_for_ais:
      z_init, _, _ = self._recog(images)
    else:
      z_init, _ = self._prior(batch=tf.shape(images)[0])

    if self._use_bijector_for_ais:
      _, _, bijector = self._recog(images)
    else:
      bijector = None
    ais_outputs = utils.AIS(ProposalLogProbFn, TargetLogProbFn, z_init, bijector=bijector)

    recons, _ = self._noise(None, ais_outputs.z_fin)

    tf.summary.image("data", utils.StitchImages(images[:64]))
    tf.summary.image("recon_means", utils.StitchImages(recons[:64]))
    tf.summary.scalar("p_accept", tf.reduce_mean(ais_outputs.p_accept))

    return AISOutputs(
        log_p=tf.reduce_logsumexp(
            tf.reshape(ais_outputs.log_p, [num_chains, -1]) - tf.log(
                tf.to_float(num_chains)), 0),
        p_accept=ais_outputs.p_accept,
        recon=recons,
        z_fin=ais_outputs.z_fin)
//  [python]
def AIS(self, images, num_chains):

    def ProposalLogProbFn(*z):
      if self._use_q_z_for_ais:
        _, log_p_z, _ = self._recog(images, z=z)
      else:
        _, log_p_z = self._prior(z)
      return tf.add_n(log_p_z)

    def TargetLogProbFn(*z):
      _, log_p_z = self._prior(z)
      _, log_p_x_z = self._noise(images, z)
      return tf.add_n(log_p_z) + log_p_x_z

    images = tf.tile(images, [num_chains, 1, 1, 1])
    if self._use_q_z_for_ais:
      z_init, _, _ = self._recog(images)
    else:
      z_init, _ = self._prior(batch=tf.shape(images)[0])

    if self._use_bijector_for_ais:
      _, _, bijector = self._recog(images)
    else:
      bijector = None
    ais_outputs = utils.AIS(ProposalLogProbFn, TargetLogProbFn, z_init, bijector=bijector)

    recons, _ = self._noise(None, ais_outputs.z_fin)

    tf.summary.image("data", utils.StitchImages(images[:64]))
    tf.summary.image("recon_means", utils.StitchImages(recons[:64]))
    tf.summary.scalar("p_accept", tf.reduce_mean(ais_outputs.p_accept))

    return AISOutputs(
        log_p=tf.reduce_logsumexp(
            tf.reshape(ais_outputs.log_p, [num_chains, -1]) - tf.log(
                tf.cast(num_chains, dtype=tf.float32)), 0),
        p_accept=ais_outputs.p_accept,
        recon=recons,
        z_fin=ais_outputs.z_fin)
//  [python]
def _predictions_for_gradients(predictions, seq_mask, delta_time,
                               attribution_max_delta_time, averaged):
  """Aggregates eligible predictions over time.

  Predictions are eligible if their are within the sequence_length (as indicated
  by seq_mask) and their associated delta_time is at most
  attribution_max_delta_time.
  Eligible predictions are either averaged across those eligble times (if
  averaged=True) or summed otherwise.

  Args:
    predictions: A Tensor of shape [batch_size, max_seq_len, 1]
      with the predictions in the sequence.
    seq_mask: A Tensor of shape [batch_size, max_sequence_length, 1] indicating
      which timesteps are padded.
    delta_time: A Tensor of shape [batch_size, max_sequence_length] describing
      the time to prediction.
    attribution_max_delta_time: Attribution is limited to values that are no
      older than that many seconds at time of prediction.
    averaged: Whether predictions are simply summed up across the time-steps
      or averaged over on the sequence length.
  Returns:
    A Tensor of shape [batch, 1, 1] of the eligible predictions
    aggregated across time.
  """
  mask = seq_mask * tf.to_float(delta_time < attribution_max_delta_time)
  predictions *= mask
  if averaged:
    predictions /= tf.reduce_sum(mask, axis=1, keepdims=True)
  return tf.reduce_sum(predictions, axis=1, keepdims=True)
//  [python]
def _predictions_for_gradients(predictions, seq_mask, delta_time,
                               attribution_max_delta_time, averaged):
  """Aggregates eligible predictions over time.

  Predictions are eligible if their are within the sequence_length (as indicated
  by seq_mask) and their associated delta_time is at most
  attribution_max_delta_time.
  Eligible predictions are either averaged across those eligble times (if
  averaged=True) or summed otherwise.

  Args:
    predictions: A Tensor of shape [batch_size, max_seq_len, 1]
      with the predictions in the sequence.
    seq_mask: A Tensor of shape [batch_size, max_sequence_length, 1] indicating
      which timesteps are padded.
    delta_time: A Tensor of shape [batch_size, max_sequence_length] describing
      the time to prediction.
    attribution_max_delta_time: Attribution is limited to values that are no
      older than that many seconds at time of prediction.
    averaged: Whether predictions are simply summed up across the time-steps
      or averaged over on the sequence length.
  Returns:
    A Tensor of shape [batch, 1, 1] of the eligible predictions
    aggregated across time.
  """
  mask = seq_mask * tf.cast(delta_time < attribution_max_delta_time, tf.float32)
  predictions *= mask
  if averaged:
    predictions /= tf.reduce_sum(mask, axis=1, keepdims=True)
  return tf.reduce_sum(predictions, axis=1, keepdims=True)
//  [python]
def AIS(self, images, num_chains):
    outputs = self.MakeVAE(images)

    def ProposalLogProbFn(*z):
      if self._use_q_z_for_ais:
        _, log_p_z = self._recog(images, z=z)
      else:
        _, log_p_z = self._prior(z)
      return tf.add_n(log_p_z)

    def TargetLogProbFn(*z):
      _, log_p_z = self._prior(z)
      _, log_p_x_z = self._noise(images, z)
      return tf.add_n(log_p_z) + log_p_x_z

    images = tf.tile(images, [num_chains, 1, 1, 1])
    if self._use_q_z_for_ais:
      z_init, _ = self._recog(images)
    else:
      z_init, _ = self._prior(batch=tf.shape(images)[0])

    ais_outputs = utils.AIS(ProposalLogProbFn, TargetLogProbFn, z_init)

    recons, _ = self._noise(None, ais_outputs.z_fin)

    tf.summary.image("data", utils.StitchImages(images[:64]))
    tf.summary.image("recon_means", utils.StitchImages(recons[:64]))
    tf.summary.scalar("p_accept", tf.reduce_mean(ais_outputs.p_accept))

    return AISOutputs(
        log_p=tf.reduce_logsumexp(
            tf.reshape(ais_outputs.log_p, [num_chains, -1]) - tf.log(
                tf.to_float(num_chains)), 0),
        p_accept=ais_outputs.p_accept,
        recon=recons,
        z_fin=ais_outputs.z_fin)
//  [python]
def AIS(self, images, num_chains):
    outputs = self.MakeVAE(images)

    def ProposalLogProbFn(*z):
      if self._use_q_z_for_ais:
        _, log_p_z = self._recog(images, z=z)
      else:
        _, log_p_z = self._prior(z)
      return tf.add_n(log_p_z)

    def TargetLogProbFn(*z):
      _, log_p_z = self._prior(z)
      _, log_p_x_z = self._noise(images, z)
      return tf.add_n(log_p_z) + log_p_x_z

    images = tf.tile(images, [num_chains, 1, 1, 1])
    if self._use_q_z_for_ais:
      z_init, _ = self._recog(images)
    else:
      z_init, _ = self._prior(batch=tf.shape(images)[0])

    ais_outputs = utils.AIS(ProposalLogProbFn, TargetLogProbFn, z_init)

    recons, _ = self._noise(None, ais_outputs.z_fin)

    tf.summary.image("data", utils.StitchImages(images[:64]))
    tf.summary.image("recon_means", utils.StitchImages(recons[:64]))
    tf.summary.scalar("p_accept", tf.reduce_mean(ais_outputs.p_accept))

    return AISOutputs(
        log_p=tf.reduce_logsumexp(
            tf.reshape(ais_outputs.log_p, [num_chains, -1]) - tf.log(
                tf.cast(num_chains, tf.float32)), 0),
        p_accept=ais_outputs.p_accept,
        recon=recons,
        z_fin=ais_outputs.z_fin)
//  [python]
def crop_or_pad(waves, length, channels):
  """Crop or pad wave to have shape [N, length, channels].

  Args:
    waves: A 3D `Tensor` of NLC format.
    length: A Python scalar. The output wave size.
    channels: Number of output waves channels.

  Returns:
    A 3D `Tensor` of NLC format with shape [N, length, channels].
  """
  waves = tf.convert_to_tensor(waves)
  batch_size = int(waves.shape[0])
  waves_shape = tf.shape(waves)

  # Force audio length.
  pad = tf.maximum(0, length - waves_shape[1])
  right_pad = tf.to_int32(tf.to_float(pad) / 2.0)
  left_pad = pad - right_pad
  waves = tf.pad(waves, [[0, 0], [left_pad, right_pad], [0, 0]])
  waves = waves[:, :length, :]

  # Force number of channels.
  num_repeats = tf.to_int32(
      tf.ceil(tf.to_float(channels) / tf.to_float(waves_shape[2])))
  waves = tf.tile(waves, [1, 1, num_repeats])[:, :, :channels]

  waves.set_shape([batch_size, length, channels])
  return waves
//  [python]
def crop_or_pad(waves, length, channels):
  """Crop or pad wave to have shape [N, length, channels].

  Args:
    waves: A 3D `Tensor` of NLC format.
    length: A Python scalar. The output wave size.
    channels: Number of output waves channels.

  Returns:
    A 3D `Tensor` of NLC format with shape [N, length, channels].
  """
  waves = tf.convert_to_tensor(waves)
  batch_size = int(waves.shape[0])
  waves_shape = tf.shape(waves)

  # Force audio length.
  pad = tf.maximum(0, length - waves_shape[1])
  right_pad = tf.cast(tf.math.ceil(tf.cast(pad, tf.float32) / 2.0), tf.int32)
  left_pad = pad - right_pad
  waves = tf.pad(waves, [[0, 0], [left_pad, right_pad], [0, 0]])
  waves = waves[:, :length, :]

  # Force number of channels.
  num_repeats = tf.cast(
      tf.math.ceil(tf.cast(channels, tf.float32) / tf.cast(waves_shape[2], tf.float32)), tf.int32)
  waves = tf.tile(waves, [1, 1, num_repeats])[:, :, :channels]

  waves.set_shape([batch_size, length, channels])
  return waves
//  [python]
def mean_squared_error(true, pred):
  """L2 distance between tensors true and pred.

  Args:
    true: the ground truth image.
    pred: the predicted image.
  Returns:
    mean squared error between ground truth and predicted image.
  """
  result = tf.reduce_sum(
      tf.squared_difference(true, pred)) / tf.to_float(tf.size(pred))
  return result
//  [python]
def mean_squared_error(true, pred):
  """L2 distance between tensors true and pred.

  Args:
    true: the ground truth image.
    pred: the predicted image.
  Returns:
    mean squared error between ground truth and predicted image.
  """
  result = tf.reduce_sum(
      tf.squared_difference(true, pred)) / tf.cast(tf.size(pred), tf.float32)
  return result
//  [python]
def convert_attribution(attribution, sequence_feature_map, seq_mask, delta_time,
                        attribution_threshold, attribution_max_delta_time,
                        prefix=''):
  """Constructs the attribution of what inputs result in a higher prediction.

  Attribution here refers to the timesteps in which the predictions (derived
  from the logits) increased. We are only interested in increases in the
  previous attribution_max_delta_time.

  Args:
    attribution: A Tensor of shape [batch, max_sequence_length, 1] computed
      using some attribution method.
    sequence_feature_map: A dictionary from name to (Sparse)Tensor.
    seq_mask: A Tensor of shape [batch_size, max_sequence_length, 1] indicating
      which timesteps are padded.
    delta_time: A Tensor of shape [batch_size, max_sequence_length] describing
      the time to prediction.
    attribution_threshold: Attribution values below this threshold will be
      dropped.
    attribution_max_delta_time: Attribution is limited to values that are no
      older than that many seconds at time of prediction.
    prefix: A string to prepend to the feature names for the attribution_dict.
  Returns:
    A dictionary from feature names to SparseTensors of
    dense_shape [batch_size, max_sequence_length, 1].
  """
  # We do not want attribution in the padding.
  attribution *= seq_mask

  # We focus on attribution in the past 12h.
  # [batch_size, max_sequence_length, 1]
  attribution *= tf.to_float(delta_time < attribution_max_delta_time)

  # We get rid of low attribution.
  attribution_indices = tf.where(attribution > attribution_threshold)
  attribution_values = tf.gather_nd(attribution, attribution_indices)

  # Now, attribution.indices indicate in the input timesteps which we should
  # attend to.
  attribution_dict = {}
  for feature, sp_feature in sequence_feature_map.items():
    # Limitation: This is not going to work for sequence feature in which
    # the third (last/token) dimension is > 1. In that case only the first
    # token would be highlighted.
    attribution_dict[prefix + feature] = tf.sparse.expand_dims(
        tf.SparseTensor(
            indices=attribution_indices,
            values=attribution_values,
            dense_shape=tf.to_int64(tf.shape(sp_feature))), axis=1)
  return attribution_dict
//  [python]
def convert_attribution(attribution, sequence_feature_map, seq_mask, delta_time,
                        attribution_threshold, attribution_max_delta_time,
                        prefix=''):
  """Constructs the attribution of what inputs result in a higher prediction.

  Attribution here refers to the timesteps in which the predictions (derived
  from the logits) increased. We are only interested in increases in the
  previous attribution_max_delta_time.

  Args:
    attribution: A Tensor of shape [batch, max_sequence_length, 1] computed
      using some attribution method.
    sequence_feature_map: A dictionary from name to (Sparse)Tensor.
    seq_mask: A Tensor of shape [batch_size, max_sequence_length, 1] indicating
      which timesteps are padded.
    delta_time: A Tensor of shape [batch_size, max_sequence_length] describing
      the time to prediction.
    attribution_threshold: Attribution values below this threshold will be
      dropped.
    attribution_max_delta_time: Attribution is limited to values that are no
      older than that many seconds at time of prediction.
    prefix: A string to prepend to the feature names for the attribution_dict.
  Returns:
    A dictionary from feature names to SparseTensors of
    dense_shape [batch_size, max_sequence_length, 1].
  """
  # We do not want attribution in the padding.
  attribution *= seq_mask

  # We focus on attribution in the past 12h.
  # [batch_size, max_sequence_length, 1]
  attribution *= tf.cast(delta_time < attribution_max_delta_time, tf.float32)

  # We get rid of low attribution.
  attribution_indices = tf.where(attribution > attribution_threshold)
  attribution_values = tf.gather_nd(attribution, attribution_indices)

  # Now, attribution.indices indicate in the input timesteps which we should
  # attend to.
  attribution_dict = {}
  for feature, sp_feature in sequence_feature_map.items():
    # Limitation: This is not going to work for sequence feature in which
    # the third (last/token) dimension is > 1. In that case only the first
    # token would be highlighted.
    attribution_dict[prefix + feature] = tf.sparse.expand_dims(
        tf.SparseTensor(
            indices=attribution_indices,
            values=attribution_values,
            dense_shape=tf.to_int64(tf.shape(sp_feature))), axis=1)
  return attribution_dict
//  [python]
def l1_error(true, pred):
  """L1 distance between tensors true and pred."""
  return tf.reduce_sum(tf.abs(true - pred)) / tf.to_float(tf.size(pred))
//  [python]
def l1_error(true, pred):
  """L1 distance between tensors true and pred."""
  return tf.reduce_sum(tf.abs(true - pred)) / tf.cast(tf.size(pred), tf.float32)
//  [python]
def get_timing_signal_1d_given_position(channels,
                                        position,
                                        min_timescale=1.0,
                                        max_timescale=1.0e4):
    """Get sinusoids of diff frequencies, with timing position given.

  Adapted from add_timing_signal_1d_given_position in
  //third_party/py/tensor2tensor/layers/common_attention.py

  Args:
    channels: scalar, size of timing embeddings to create. The number of
        different timescales is equal to channels / 2.
    position: a Tensor with shape [batch, seq_len]
    min_timescale: a float
    max_timescale: a float

  Returns:
    a Tensor of timing signals [batch, seq_len, channels]
  """
    num_timescales = channels // 2
    log_timescale_increment = (
            math.log(float(max_timescale) / float(min_timescale)) /
            (tf.to_float(num_timescales) - 1))
    inv_timescales = min_timescale * tf.exp(
        tf.to_float(tf.range(num_timescales)) * -log_timescale_increment)
    scaled_time = (
            tf.expand_dims(tf.to_float(position), 2) * tf.expand_dims(
        tf.expand_dims(inv_timescales, 0), 0))
    signal = tf.concat([tf.sin(scaled_time), tf.cos(scaled_time)], axis=2)
    signal = tf.pad(signal, [[0, 0], [0, 0], [0, tf.mod(channels, 2)]])
    return signal
//  [python]
def get_timing_signal_1d_given_position(channels,
                                        position,
                                        min_timescale=1.0,
                                        max_timescale=1.0e4):
    """Get sinusoids of diff frequencies, with timing position given.

  Adapted from add_timing_signal_1d_given_position in
  //third_party/py/tensor2tensor/layers/common_attention.py

  Args:
    channels: scalar, size of timing embeddings to create. The number of
        different timescales is equal to channels / 2.
    position: a Tensor with shape [batch, seq_len]
    min_timescale: a float
    max_timescale: a float

  Returns:
    a Tensor of timing signals [batch, seq_len, channels]
  """
    num_timescales = channels // 2
    log_timescale_increment = (
            math.log(float(max_timescale) / float(min_timescale)) /
            (tf.cast(num_timescales, tf.float32) - 1))
    inv_timescales = min_timescale * tf.exp(
        tf.cast(tf.range(num_timescales), tf.float32) * -log_timescale_increment)
    scaled_time = (
            tf.expand_dims(tf.cast(position, tf.float32), 2) * tf.expand_dims(
        tf.expand_dims(inv_timescales, 0), 0))
    signal = tf.concat([tf.sin(scaled_time), tf.cos(scaled_time)], axis=2)
    signal = tf.pad(signal, [[0, 0], [0, 0], [0, tf.mod(channels, 2)]])
    return signal
//  [python]
def get_rnn_text_classification_family(cfg):
  """Get a task for the given cfg.

  Args:
    cfg: config specifying the model generated by
      `sample_rnn_text_classification_family_cfg`.

  Returns:
    base.BaseTask for the given config.
  """

  w_init = utils.get_initializer(cfg["w_init"])
  init = {"w": w_init}

  def _build(batch):
    """Build the sonnet module.

    Args:
      batch: A dictionary with keys "label", "label_onehot", and "text" mapping
        to tensors. The "text" consists of int tokens. These tokens are
        truncated to the length of the vocab before performing an embedding
        lookup.

    Returns:
      Loss of the batch.
    """
    vocab_size = cfg["vocab_size"]
    max_token = cfg["dataset"][1]["max_token"]
    if max_token:
      vocab_size = min(max_token, vocab_size)

    # Clip the max token to be vocab_size-1.
    tokens = tf.minimum(
        tf.to_int32(batch["text"]),
        tf.to_int32(tf.reshape(vocab_size - 1, [1, 1])))

    embed = snt.Embed(vocab_size=vocab_size, embed_dim=cfg["embed_dim"])
    embedded_tokens = embed(tokens)
    rnn = utils.get_rnn_core(cfg["core"])

    batch_size = tokens.shape.as_list()[0]

    state = rnn.initial_state(batch_size, trainable=cfg["trainable_init"])

    outputs, _ = tf.nn.dynamic_rnn(rnn, embedded_tokens, initial_state=state)
    if cfg["loss_compute"] == "last":
      rnn_output = outputs[:, -1]  # grab the last output
    elif cfg["loss_compute"] == "avg":
      rnn_output = tf.reduce_mean(outputs, 1)  # average over length
    elif cfg["loss_compute"] == "max":
      rnn_output = tf.reduce_max(outputs, 1)
    else:
      raise ValueError("Not supported loss_compute [%s]" % cfg["loss_compute"])

    logits = snt.Linear(
        batch["label_onehot"].shape[1], initializers=init)(
            rnn_output)

    loss_vec = tf.nn.softmax_cross_entropy_with_logits_v2(
        labels=batch["label_onehot"], logits=logits)
    return tf.reduce_mean(loss_vec)

  datasets = utils.get_text_dataset(cfg["dataset"])
  return base.DatasetModelTask(lambda: snt.Module(_build), datasets)
//  [python]
def get_rnn_text_classification_family(cfg):
  """Get a task for the given cfg.

  Args:
    cfg: config specifying the model generated by
      `sample_rnn_text_classification_family_cfg`.

  Returns:
    base.BaseTask for the given config.
  """

  w_init = utils.get_initializer(cfg["w_init"])
  init = {"w": w_init}

  def _build(batch):
    """Build the sonnet module.

    Args:
      batch: A dictionary with keys "label", "label_onehot", and "text" mapping
        to tensors. The "text" consists of int tokens. These tokens are
        truncated to the length of the vocab before performing an embedding
        lookup.

    Returns:
      Loss of the batch.
    """
    vocab_size = cfg["vocab_size"]
    max_token = cfg["dataset"][1]["max_token"]
    if max_token:
      vocab_size = min(max_token, vocab_size)

    # Clip the max token to be vocab_size-1.
    tokens = tf.minimum(
        tf.cast(batch["text"], tf.int32),
        tf.cast(tf.reshape(vocab_size - 1, [1, 1]), tf.int32))

    embed = snt.Embed(vocab_size=vocab_size, embed_dim=cfg["embed_dim"])
    embedded_tokens = embed(tokens)
    rnn = utils.get_rnn_core(cfg["core"])

    batch_size = tokens.shape.as_list()[0]

    state = rnn.initial_state(batch_size, trainable=cfg["trainable_init"])

    outputs, _ = tf.nn.dynamic_rnn(rnn, embedded_tokens, initial_state=state)
    if cfg["loss_compute"] == "last":
      rnn_output = outputs[:, -1]  # grab the last output
    elif cfg["loss_compute"] == "avg":
      rnn_output = tf.reduce_mean(outputs, 1)  # average over length
    elif cfg["loss_compute"] == "max":
      rnn_output = tf.reduce_max(outputs, 1)
    else:
      raise ValueError("Not supported loss_compute [%s]" % cfg["loss_compute"])

    logits = snt.Linear(
        batch["label_onehot"].shape[1], initializers=init)(
            rnn_output)

    loss_vec = tf.nn.softmax_cross_entropy_with_logits_v2(
        labels=batch["label_onehot"], logits=logits)
    return tf.reduce_mean(loss_vec)

  datasets = utils.get_text_dataset(cfg["dataset"])
  return base.DatasetModelTask(lambda: snt.Module(_build), datasets)
//  [python]
def soft_cross_entropy_loss(logits, labels, smoothing, vocab_size):
  """Calculate cross entropy loss while ignoring padding.
  Args:
    logits: Tensor of size [batch_size, length_logits, vocab_size]
    labels: Tensor of size [batch_size, length_labels]
    smoothing: Label smoothing constant, used to determine the on and off values
    vocab_size: int size of the vocabulary
  Returns:
    Returns the cross entropy loss and weight tensors: float32 tensors with
      shape [batch_size, max(length_logits, length_labels)]
  """
  with tf.name_scope("loss", values=[logits, labels]):
    # Calculate smoothing cross entropy
    with tf.name_scope("smoothing_cross_entropy", values=[logits, labels]):
      confidence = 1.0 - smoothing
      low_confidence = (1.0 - confidence) / tf.to_float(vocab_size - 1)
      soft_targets = tf.one_hot(
          tf.cast(labels, tf.int32),
          depth=vocab_size,
          on_value=confidence,
          off_value=low_confidence)
      xentropy = tf.nn.softmax_cross_entropy_with_logits_v2(
          logits=logits, labels=soft_targets)

      # Calculate the best (lowest) possible value of cross entropy, and
      # subtract from the cross entropy loss.
      normalizing_constant = -(
          confidence * tf.log(confidence) + tf.to_float(vocab_size - 1) *
          low_confidence * tf.log(low_confidence + 1e-20))
      xentropy -= normalizing_constant
    return xentropy
//  [python]
def soft_cross_entropy_loss(logits, labels, smoothing, vocab_size):
  """Calculate cross entropy loss while ignoring padding.
  Args:
    logits: Tensor of size [batch_size, length_logits, vocab_size]
    labels: Tensor of size [batch_size, length_labels]
    smoothing: Label smoothing constant, used to determine the on and off values
    vocab_size: int size of the vocabulary
  Returns:
    Returns the cross entropy loss and weight tensors: float32 tensors with
      shape [batch_size, max(length_logits, length_labels)]
  """
  with tf.name_scope("loss", values=[logits, labels]):
    # Calculate smoothing cross entropy
    with tf.name_scope("smoothing_cross_entropy", values=[logits, labels]):
      confidence = 1.0 - smoothing
      low_confidence = (1.0 - confidence) / tf.cast(vocab_size - 1, tf.float32)
      soft_targets = tf.one_hot(
          tf.cast(labels, tf.int32),
          depth=vocab_size,
          on_value=confidence,
          off_value=low_confidence)
      xentropy = tf.nn.softmax_cross_entropy_with_logits_v2(
          logits=logits, labels=soft_targets)

      # Calculate the best (lowest) possible value of cross entropy, and
      # subtract from the cross entropy loss.
      normalizing_constant = -(
          confidence * tf.log(confidence) + tf.cast(vocab_size - 1, tf.float32) *
          low_confidence * tf.log(low_confidence + 1e-20))
      xentropy -= normalizing_constant
    return xentropy