jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def forward(self, examples, labels):
    """Build the graph for the forward pass."""
    opts = self._options

    # Declare all variables we need.
    # Embedding: [vocab_size, emb_dim]
    init_width = 0.5 / opts.emb_dim
    emb = tf.Variable(
        tf.random_uniform(
            [opts.vocab_size, opts.emb_dim], -init_width, init_width),
        name="emb")
    self._emb = emb

    # Softmax weight: [vocab_size, emb_dim]. Transposed.
    sm_w_t = tf.Variable(
        tf.zeros([opts.vocab_size, opts.emb_dim]),
        name="sm_w_t")

    # Softmax bias: [vocab_size].
    sm_b = tf.Variable(tf.zeros([opts.vocab_size]), name="sm_b")

    # Global step: scalar, i.e., shape [].
    self.global_step = tf.Variable(0, name="global_step")

    # Nodes to compute the nce loss w/ candidate sampling.
    labels_matrix = tf.reshape(
        tf.cast(labels,
                dtype=tf.int64),
        [opts.batch_size, 1])

    # Negative sampling.
    sampled_ids, _, _ = (tf.nn.fixed_unigram_candidate_sampler(
        true_classes=labels_matrix,
        num_true=1,
        num_sampled=opts.num_samples,
        unique=True,
        range_max=opts.vocab_size,
        distortion=0.75,
        unigrams=opts.vocab_counts.tolist()))

    # Embeddings for examples: [batch_size, emb_dim]
    example_emb = tf.nn.embedding_lookup(emb, examples)

    # Weights for labels: [batch_size, emb_dim]
    true_w = tf.nn.embedding_lookup(sm_w_t, labels)
    # Biases for labels: [batch_size, 1]
    true_b = tf.nn.embedding_lookup(sm_b, labels)

    # Weights for sampled ids: [num_sampled, emb_dim]
    sampled_w = tf.nn.embedding_lookup(sm_w_t, sampled_ids)
    # Biases for sampled ids: [num_sampled, 1]
    sampled_b = tf.nn.embedding_lookup(sm_b, sampled_ids)

    # True logits: [batch_size, 1]
    true_logits = tf.reduce_sum(tf.multiply(example_emb, true_w), 1) + true_b

    # Sampled logits: [batch_size, num_sampled]
    # We replicate sampled noise labels for all examples in the batch
    # using the matmul.
    sampled_b_vec = tf.reshape(sampled_b, [opts.num_samples])
    sampled_logits = tf.matmul(example_emb,
                               sampled_w,
                               transpose_b=True) + sampled_b_vec
    return true_logits, sampled_logits
//  [python]
def aggregate_loss_top_k(self, pixelwise_loss, num_valid_pixels=None):
    """Aggregate the top-k greatest pixelwise loss.

    Args:
      pixelwise_loss: a float tensor in shape (batch_size, height, width) which
        stores the loss of each pixel.
      num_valid_pixels: the number of pixels which are not ignored. If None, all
        the pixels are valid.

    Returns:
       A 0-D float which stores the overall loss of the batch.
    """
    pixelwise_loss = tf.reshape(pixelwise_loss, shape=[-1])
    top_k_pixels = tf.to_int32(
        self._top_k_percent_pixels
        * tf.to_float(tf.size(pixelwise_loss)),
    )
    top_k_losses, _ = tf.math.top_k(pixelwise_loss, k=top_k_pixels)
    normalizer = tf.to_float(top_k_pixels)
    if num_valid_pixels is not None:
      normalizer = tf.minimum(normalizer,
                              tf.to_float(num_valid_pixels))
    return tf.reduce_sum(top_k_losses) / (normalizer + EPSILON)
//  [python]
def aggregate_loss_top_k(self, pixelwise_loss, num_valid_pixels=None):
    """Aggregate the top-k greatest pixelwise loss.

    Args:
      pixelwise_loss: a float tensor in shape (batch_size, height, width) which
        stores the loss of each pixel.
      num_valid_pixels: the number of pixels which are not ignored. If None, all
        the pixels are valid.

    Returns:
       A 0-D float which stores the overall loss of the batch.
    """
    pixelwise_loss = tf.reshape(pixelwise_loss, shape=[-1])
    top_k_pixels = tf.cast(
        self._top_k_percent_pixels
        * tf.cast(tf.size(pixelwise_loss), tf.float32),
        tf.int32,
    )
    top_k_losses, _ = tf.math.top_k(pixelwise_loss, k=top_k_pixels)
    normalizer = tf.cast(top_k_pixels, top_k_losses.dtype)
    if num_valid_pixels is not None:
      normalizer = tf.minimum(normalizer,
                              tf.cast(num_valid_pixels, top_k_losses.dtype))
    return tf.reduce_sum(top_k_losses) / (normalizer + EPSILON)
//  [python]
def test_mask_invalid_locations(self):
    hidden_states = self._get_hidden_states()
    batch_size = 1
    seq_length = 8
    hidden_size = 4
    hidden_states = tf.reshape(hidden_states,
                               (batch_size, seq_length, hidden_size))
    hidden_states = longformer_attention.LongformerAttention._chunk(
        hidden_states, window_overlap=2)

    hid_states_1 = longformer_attention.LongformerAttention._mask_invalid_locations(
        hidden_states, 1)
    hid_states_2 = longformer_attention.LongformerAttention._mask_invalid_locations(
        hidden_states, 2)
    hid_states_3 = longformer_attention.LongformerAttention._mask_invalid_locations(
        hidden_states[:, :, :, :3], 2)
    hid_states_4 = longformer_attention.LongformerAttention._mask_invalid_locations(
        hidden_states[:, :, 2:, :], 2)

    self.assertEqual(
        tf.math.reduce_sum(
            tf.cast(tf.math.is_inf(hid_states_1), tf.int32)), 8)
    self.assertEqual(
        tf.math.reduce_sum(
            tf.cast(tf.math.is_inf(hid_states_2), tf.int32)), 24)
    self.assertEqual(
        tf.math.reduce_sum(
            tf.cast(tf.math.is_inf(hid_states_3), tf.int32)), 24)
    self.assertEqual(
        tf.math.reduce_sum(
            tf.cast(tf.math.is_inf(hid_states_4), tf.int32)), 12)
//  [python]
def test_mask_invalid_locations(self):
    hidden_states = self._get_hidden_states()
    batch_size = 1
    seq_length = 8
    hidden_size = 4
    hidden_states = tf.reshape(hidden_states,
                               (batch_size, seq_length, hidden_size))
    hidden_states = longformer_attention.LongformerAttention._chunk(
        hidden_states, window_overlap=2)

    hid_states_1 = longformer_attention.LongformerAttention._mask_invalid_locations(
        hidden_states, 1)
    hid_states_2 = longformer_attention.LongformerAttention._mask_invalid_locations(
        hidden_states, 2)
    hid_states_3 = longformer_attention.LongformerAttention._mask_invalid_locations(
        hidden_states[:, :, :, :3], 2)
    hid_states_4 = longformer_attention.LongformerAttention._mask_invalid_locations(
        hidden_states[:, :, 2:, :], 2)

    self.assertEqual(
        tf.math.reduce_sum(
            tf.cast(tf.math.is_inf(hid_states_1), tf.dtypes.int32)), 8)
    self.assertEqual(
        tf.math.reduce_sum(
            tf.cast(tf.math.is_inf(hid_states_2), tf.dtypes.int32)), 24)
    self.assertEqual(
        tf.math.reduce_sum(
            tf.cast(tf.math.is_inf(hid_states_3), tf.dtypes.int32)), 24)
    self.assertEqual(
        tf.math.reduce_sum(
            tf.cast(tf.math.is_inf(hid_states_4), tf.dtypes.int32)), 12)
//  [python]
def get_padding(x, padding_value=0):
  """Return float tensor representing the padding values in x.
  Args:
    x: int tensor with any shape
    padding_value: int value that
  Returns:
    float tensor with same shape as x containing values 0 or 1.
      0 -> non-padding, 1 -> padding
  """
  with tf.name_scope("padding"):
    return tf.to_float(tf.equal(x, padding_value))
//  [python]
def get_padding(x, padding_value=0):
  """Return float tensor representing the padding values in x.
  Args:
    x: int tensor with any shape
    padding_value: int value that
  Returns:
    float tensor with same shape as x containing values 0 or 1.
      0 -> non-padding, 1 -> padding
  """
  with tf.name_scope("padding"):
    return tf.cast(tf.equal(x, padding_value), tf.float32)
//  [python]
def prepare_centre_neighbour_sentences(item, window_length):

    n_sentences = t.cast(item.row_splits[1], t.int32)
    # 1 - n_sentences-1
    centre_index = tf.random.uniform(minval=1, maxval=n_sentences - 1, shape=(), dtype=tf.int32)

    min_window_boundary = tf.maximum(0, centre_index - window_length)
    max_window_boundary = tf.minimum(centre_index + window_length, n_sentences)

    neighbour_indexes = tf.squeeze(
        tf.where(tf.not_equal(tf.range(min_window_boundary, max_window_boundary), centre_index)), axis=1
    )

    # Randomly sample it for each example with different number of neighbours
    # n_neighbours = tf.random.uniform(minval=1, maxval = tf.shape(neighbour_indexes)[0], shape=(), dtype=tf.int32)
    # neighbour_index_ids = tf.random.uniform(minval=0, \
    # maxval = tf.shape(neighbour_indexes)[0], shape=(n_neighbours,), dtype=tf.int32)
    # neighbour_indexes = tf.gather(neighbour_indexes, neighbour_index_ids)

    centre_sentence = tf.gather(item[0], centre_index)
    neighbour_sentence = tf.gather(item[0], neighbour_indexes)
    # Repeat n-times
    n = tf.shape(neighbour_sentence)[0]
    centre_sentence = tf.repeat(centre_sentence, [n])

    # vec = tf.expand_dims(tf.range(batch_size), 0)

    # We use this same vector for all similiar sentences
    # Then we calculate scores (M x M transpose)
    # use this scores to generate mask for softmax
    # vec = tf.stop_gradient(tf.random.uniform(shape=(1, 32))) # keep a unique random
    # multiply = tf.stack([n, 1], axis=0)
    # tensor = tf.cast(tf.tile(vec, multiply), tf.float32)

    return {'centre_sentence': centre_sentence, 'neighbour_sentence': neighbour_sentence}
//  [python]
def prepare_centre_neighbour_sentences(item, window_length):

    n_sentences = tf.cast(item.row_splits[1], tf.int32)
    # 1 - n_sentences-1
    centre_index = tf.random.uniform(minval=1, maxval=n_sentences - 1, shape=(), dtype=tf.int32)

    min_window_boundary = tf.maximum(0, centre_index - window_length)
    max_window_boundary = tf.minimum(centre_index + window_length, n_sentences)

    neighbour_indexes = tf.squeeze(
        tf.where(tf.not_equal(tf.range(min_window_boundary, max_window_boundary), centre_index)), axis=1
    )

    # Randomly sample it for each example with different number of neighbours
    # n_neighbours = tf.random.uniform(minval=1, maxval = tf.shape(neighbour_indexes)[0], shape=(), dtype=tf.int32)
    # neighbour_index_ids = tf.random.uniform(minval=0, \
    # maxval = tf.shape(neighbour_indexes)[0], shape=(n_neighbours,), dtype=tf.int32)
    # neighbour_indexes = tf.gather(neighbour_indexes, neighbour_index_ids)

    centre_sentence = tf.gather(item[0], centre_index)
    neighbour_sentence = tf.gather(item[0], neighbour_indexes)
    # Repeat n-times
    n = tf.shape(neighbour_sentence)[0]
    centre_sentence = tf.repeat(centre_sentence, [n])

    # vec = tf.expand_dims(tf.range(batch_size), 0)

    # We use this same vector for all similiar sentences
    # Then we calculate scores (M x M transpose)
    # use this scores to generate mask for softmax
    # vec = tf.stop_gradient(tf.random.uniform(shape=(1, 32))) # keep a unique random
    # multiply = tf.stack([n, 1], axis=0)
    # tensor = tf.cast(tf.tile(vec, multiply), tf.float32)

    return {'centre_sentence': centre_sentence, 'neighbour_sentence': neighbour_sentence}
//  [python]
def get_position_encoding(
    length, hidden_size, min_timescale=1.0, max_timescale=1.0e4):
  """Return positional encoding.
  Calculates the position encoding as a mix of sine and cosine functions with
  geometrically increasing wavelengths.
  Defined and formulized in Attention is All You Need, section 3.5.
  Args:
    length: Sequence length.
    hidden_size: Size of the
    min_timescale: Minimum scale that will be applied at each position
    max_timescale: Maximum scale that will be applied at each position
  Returns:
    Tensor with shape [length, hidden_size]
  """
  position = t.cast(t.range(length), t.float32)
  num_timescales = hidden_size // 2
  log_timescale_increment = (
      math.log(float(max_timescale) / float(min_timescale)) /
      (t.cast(num_timescales, t.float32) - 1))
  inv_timescales = min_timescale * t.exp(
      t.cast(t.range(num_timescales), t.float32) * -log_timescale_increment)
  scaled_time = t.expand_dims(position, 1) * t.expand_dims(inv_timescales, 0)
  signal = t.concat([t.sin(scaled_time), t.cos(scaled_time)], axis=1)
  return signal
//  [python]
def get_position_encoding(
    length, hidden_size, min_timescale=1.0, max_timescale=1.0e4):
  """Return positional encoding.
  Calculates the position encoding as a mix of sine and cosine functions with
  geometrically increasing wavelengths.
  Defined and formulized in Attention is All You Need, section 3.5.
  Args:
    length: Sequence length.
    hidden_size: Size of the
    min_timescale: Minimum scale that will be applied at each position
    max_timescale: Maximum scale that will be applied at each position
  Returns:
    Tensor with shape [length, hidden_size]
  """
  position = tf.cast(tf.range(length), tf.float32)
  num_timescales = hidden_size // 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(position, 1) * tf.expand_dims(inv_timescales, 0)
  signal = tf.concat([tf.sin(scaled_time), tf.cos(scaled_time)], axis=1)
  return signal
//  [python]
def get_actual_mask_scores(logits, labels, ignore_label):
  """Gets actual mask scores."""
  _, height, width, num_classes = logits.get_shape().as_list()
  batch_size = tf.shape(logits)[0]
  logits = tf.stop_gradient(logits)
  labels = tf.image.resize(
      labels, (height, width), method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)
  predicted_labels = tf.argmax(logits, -1, output_type=tf.int32)
  flat_predictions = tf.reshape(predicted_labels, [batch_size, -1])
  flat_labels = tf.cast(tf.reshape(labels, [batch_size, -1]), tf.int32)

  one_hot_predictions = tf.one_hot(
      flat_predictions, num_classes, on_value=True, off_value=False)
  one_hot_labels = tf.one_hot(
      flat_labels, num_classes, on_value=True, off_value=False)
  keep_mask = tf.not_equal(flat_labels, ignore_label)
  keep_mask = tf.expand_dims(keep_mask, 2)

  overlap = tf.logical_and(one_hot_predictions, one_hot_labels)
  overlap = tf.logical_and(overlap, keep_mask)
  overlap = tf.reduce_sum(tf.to_float(overlap), axis=1)
  union = tf.logical_or(one_hot_predictions, one_hot_labels)
  union = tf.logical_and(union, keep_mask)
  union = tf.reduce_sum(tf.to_float(union), axis=1)
  actual_scores = tf.divide(overlap, tf.maximum(union, EPSILON))
  return actual_scores
//  [python]
def get_actual_mask_scores(logits, labels, ignore_label):
  """Gets actual mask scores."""
  _, height, width, num_classes = logits.get_shape().as_list()
  batch_size = tf.shape(logits)[0]
  logits = tf.stop_gradient(logits)
  labels = tf.image.resize(
      labels, (height, width), method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)
  predicted_labels = tf.argmax(logits, -1, output_type=tf.int32)
  flat_predictions = tf.reshape(predicted_labels, [batch_size, -1])
  flat_labels = tf.cast(tf.reshape(labels, [batch_size, -1]), tf.int32)

  one_hot_predictions = tf.one_hot(
      flat_predictions, num_classes, on_value=True, off_value=False)
  one_hot_labels = tf.one_hot(
      flat_labels, num_classes, on_value=True, off_value=False)
  keep_mask = tf.not_equal(flat_labels, ignore_label)
  keep_mask = tf.expand_dims(keep_mask, 2)

  overlap = tf.logical_and(one_hot_predictions, one_hot_labels)
  overlap = tf.logical_and(overlap, keep_mask)
  overlap = tf.reduce_sum(tf.cast(overlap, tf.float32), axis=1)
  union = tf.logical_or(one_hot_predictions, one_hot_labels)
  union = tf.logical_and(union, keep_mask)
  union = tf.reduce_sum(tf.cast(union, tf.float32), axis=1)
  actual_scores = tf.divide(overlap, tf.maximum(union, EPSILON))
  return actual_scores
//  [python]
def loss_fn(ytrue, ypred):
    scce = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
    mask = tf.math.logical_not(tf.math.equal(ytrue, 0))
    mask = tf.to_int64(mask)
    loss = scce(ytrue, ypred, sample_weight=mask)
    return loss
//  [python]
def loss_fn(ytrue, ypred):
    scce = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
    mask = tf.math.logical_not(tf.math.equal(ytrue, 0))
    mask = tf.cast(mask, dtype=tf.int64)
    loss = scce(ytrue, ypred, sample_weight=mask)
    return loss
//  [python]
def get_labels_with_prob(self, logits, labels, valid_mask, **unused_kwargs):
    """Get a tensor representing the probability of each class for each pixel.

    This method can be overridden in subclasses for customizing loss function.

    Args:
      logits: A float tensor in shape (batch_size, height, width, num_classes)
        which is the output of the network.
      labels: A tensor in shape (batch_size, height, width, num_layers), which
        is the label masks of the ground truth. The num_layers can be > 1 if the
        pixels are labeled as multiple classes.
      valid_mask: A bool tensor in shape (batch_size, height, width, num_layers)
        which indicates the ignored labels in each ground truth layer.
      **unused_kwargs: Unused keyword arguments.

    Returns:
       A float tensor in shape (batch_size, height, width, num_classes).
    """
    num_classes = logits.get_shape().as_list()[-1]

    if self._gt_is_matting_map:
      # (batch_size, height, width, num_classes=2)
      train_labels = tf.concat([1 - labels, labels], axis=-1)
    else:
      labels = tf.cast(labels, tf.int32)
      # Assign pixel with ignore label to class -1, which will be ignored by
      # tf.one_hot operation.
      # (batch_size, height, width, num_masks)
      labels = tf.where(valid_mask, labels, -tf.ones_like(labels))

      if self._use_binary_cross_entropy:
        # (batch_size, height, width, num_masks, num_classes)
        one_hot_labels_per_mask = tf.one_hot(
            labels,
            depth=num_classes,
            on_value=True,
            off_value=False,
            dtype=tf.bool,
            axis=-1)
        # Aggregate all one-hot labels to get a binary mask in shape
        # (batch_size, height, width, num_classes), which represents all the
        # classes that a pixel is labeled as.
        # For example, if a pixel is labeled as "window" (id=1) and also being a
        # part of the "building" (id=3), then its train_labels are [0,1,0,1].
        train_labels = tf.cast(
            tf.reduce_any(one_hot_labels_per_mask, axis=-2), dtype=logits.dtype)
      else:
        # (batch_size, height, width, num_classes)
        train_labels = tf.one_hot(
            tf.squeeze(labels, axis=-1), depth=num_classes, dtype=logits.dtype)

    return train_labels * (
        1 - self._label_smoothing) + self._label_smoothing / num_classes
//  [python]
def call(self, q, k, v, mask):
        """This is where the layer's logic lives."""
        matmul_qk = tf.matmul(q, k, transpose_b=True)  # (..., seq_len_q, seq_len_k)

        # scale matmul_qk
        dk = tf.to_float(tf.shape(k)[-1])
        scaled_attention_logits = matmul_qk / tf.math.sqrt(dk)

        if self.uni:
            uni_mask = tf.ones(tf.shape(scaled_attention_logits))
            uni_mask = tf.linalg.band_part(uni_mask, -1, self.look_ahead)
            scaled_attention_logits += (1 - uni_mask) * -1e9
        # add the mask to the scaled tensor.
        if mask is not None:
            # scaled_attention_logits += mask * -float('inf') #-1e9
            scaled_attention_logits = tf.where(tf.to_bool(mask), -float('inf'), scaled_attention_logits)
        # softmax is normalized on the last axis (seq_len_k) so that the scores
        # add up to 1.
        # (..., seq_len_q, seq_len_k)
        attention_weights = tf.nn.softmax(scaled_attention_logits, axis=-1)
        attention_weights = tf.where(tf.to_bool(mask), 0.0, attention_weights)
        output = tf.matmul(attention_weights, v)  # (..., seq_len_q, depth_v)

        return output, attention_weights
//  [python]
def call(self, q, k, v, mask):
        """This is where the layer's logic lives."""
        matmul_qk = tf.matmul(q, k, transpose_b=True)  # (..., seq_len_q, seq_len_k)

        # scale matmul_qk
        dk = tf.cast(tf.shape(k)[-1], tf.float32)
        scaled_attention_logits = matmul_qk / tf.math.sqrt(dk)

        if self.uni:
            uni_mask = tf.ones(tf.shape(scaled_attention_logits))
            uni_mask = tf.linalg.band_part(uni_mask, -1, self.look_ahead)
            scaled_attention_logits += (1 - uni_mask) * -1e9
        # add the mask to the scaled tensor.
        if mask is not None:
            # scaled_attention_logits += mask * -float('inf') #-1e9
            scaled_attention_logits = tf.where(tf.cast(mask, tf.bool), -float('inf'), scaled_attention_logits)
        # softmax is normalized on the last axis (seq_len_k) so that the scores
        # add up to 1.
        # (..., seq_len_q, seq_len_k)
        attention_weights = tf.nn.softmax(scaled_attention_logits, axis=-1)
        attention_weights = tf.where(tf.cast(mask, tf.bool), 0.0, attention_weights)
        output = tf.matmul(attention_weights, v)  # (..., seq_len_q, depth_v)

        return output, attention_weights
//  [python]
def update_state(
      self,
      values: tf.Tensor,
      is_treatment: tf.Tensor,
      sample_weight: Optional[tf.Tensor] = None,
      **kwargs,
  ):
    """Computes aggregate, control and treatment metric updates.

    Args:
      values: A `tf.Tensor` instance of shape (D0, ..., DN) passed to the
        `update_state` method of the treatment, control, and overall metrics.
      is_treatment: a `tf.Tensor` of shape (D0,) or (D0, 1) castable to boolean
        indicating if the example belongs to the treatment group (True) or
        control group (False).
      sample_weight: optional `tf.Tensor` for the sample weight. If given it
        will also be sliced by the is_treatment tensor.
      **kwargs: Keyword arguments that will be passed to the `update_state`
        method of each metric.
    """
    slicing_feature = tf.reshape(tf.to_float(is_treatment), [-1])
    if sample_weight is not None:
      sample_weight = tf.reshape(sample_weight, [-1])
    super().update_state(
        values,
        sample_weight=sample_weight,
        slicing_feature=slicing_feature,
        **kwargs,
    )
//  [python]
def update_state(
      self,
      values: tf.Tensor,
      is_treatment: tf.Tensor,
      sample_weight: Optional[tf.Tensor] = None,
      **kwargs,
  ):
    """Computes aggregate, control and treatment metric updates.

    Args:
      values: A `tf.Tensor` instance of shape (D0, ..., DN) passed to the
        `update_state` method of the treatment, control, and overall metrics.
      is_treatment: a `tf.Tensor` of shape (D0,) or (D0, 1) castable to boolean
        indicating if the example belongs to the treatment group (True) or
        control group (False).
      sample_weight: optional `tf.Tensor` for the sample weight. If given it
        will also be sliced by the is_treatment tensor.
      **kwargs: Keyword arguments that will be passed to the `update_state`
        method of each metric.
    """
    slicing_feature = tf.reshape(tf.cast(is_treatment, tf.bool), [-1])
    if sample_weight is not None:
      sample_weight = tf.reshape(sample_weight, [-1])
    super().update_state(
        values,
        sample_weight=sample_weight,
        slicing_feature=slicing_feature,
        **kwargs,
    )
//  [python]
def call(self, y_true, y_pred):
        target_threshold = t.cast(self.target_threshold, y_true.dtype)
        y_true = tf.where(y_true > target_threshold, tf.ones_like(y_true), tf.zeros_like(y_true))
        return super().call(y_true, y_pred)
//  [python]
def call(self, y_true, y_pred):
        target_threshold = tf.cast(self.target_threshold, y_true.dtype)
        y_true = tf.where(y_true > target_threshold, tf.ones_like(y_true), tf.zeros_like(y_true))
        return super().call(y_true, y_pred)
//  [python]
def _compute_dloss_by_dx(op, grad):
    x = t.cast(op.inputs[0], t.float32)
    encoding_min = t.cast(op.inputs[int(QuantizeOpIndices.encoding_min)], t.float32)
    encoding_max = t.cast(op.inputs[int(QuantizeOpIndices.encoding_max)], t.float32)
    op_mode = t.cast(op.inputs[int(QuantizeOpIndices.op_mode)], t.int8)

    inner_cond = tf.compat.v2.where(tf.less_equal(x, encoding_max),  # condition to check per value
                                    1.0,  # execute if true
                                    0.0)  # execute if false

    dloss_by_dx = (tf.compat.v2.where(tf.less_equal(encoding_min, x),  # condition to check per value
                                      inner_cond,  # execute if true
                                      0.0)) * grad

    # Pass through gradient for skipped ops
    dloss_by_dx = tf.cond(tf.equal(op_mode, 3), lambda: grad, lambda: dloss_by_dx)

    return dloss_by_dx
//  [python]
def _compute_dloss_by_dx(op, grad):
    x = tf.cast(op.inputs[0], tf.float32)
    encoding_min = tf.cast(op.inputs[int(QuantizeOpIndices.encoding_min)], tf.float32)
    encoding_max = tf.cast(op.inputs[int(QuantizeOpIndices.encoding_max)], tf.float32)
    op_mode = tf.cast(op.inputs[int(QuantizeOpIndices.op_mode)], tf.int8)

    inner_cond = tf.compat.v2.where(tf.less_equal(x, encoding_max),  # condition to check per value
                                    1.0,  # execute if true
                                    0.0)  # execute if false

    dloss_by_dx = (tf.compat.v2.where(tf.less_equal(encoding_min, x),  # condition to check per value
                                      inner_cond,  # execute if true
                                      0.0)) * grad

    # Pass through gradient for skipped ops
    dloss_by_dx = tf.cond(tf.equal(op_mode, 3), lambda: grad, lambda: dloss_by_dx)

    return dloss_by_dx
//  [python]
def main(strategy):
  """Trains an MNIST model using the given tf.distribute.Strategy."""
  # TODO(fengliuai): put this in some automatically generated code.
  os.environ[
      'TF_MLIR_TFR_LIB_DIR'] = 'tensorflow/compiler/mlir/tfr/examples/mnist'

  ds_train = tfds.load('mnist', split='train', shuffle_files=True)
  ds_train = ds_train.shuffle(1024).batch(batch_size).prefetch(64)
  ds_train = strategy.experimental_distribute_dataset(ds_train)

  with strategy.scope():
    # Create an mnist float model with the specified float state.
    model = FloatModel()
    optimizer = tf.keras.optimizers.Adam(learning_rate=learning_rate)

  def train_step(features):
    inputs = tf.image.convert_image_dtype(
        features['image'], dtype=tf.float32, saturate=False)
    labels = tf.one_hot(features['label'], num_classes)

    with tf.GradientTape() as tape:
      logits = model(inputs)
      loss_value = tf.reduce_mean(
          tf.nn.softmax_cross_entropy_with_logits(labels, logits))

    grads = tape.gradient(loss_value, model.trainable_variables)
    correct_prediction = tf.equal(tf.argmax(logits, 1), tf.argmax(labels, 1))
    accuracy = tf.reduce_mean(tf.to_float(correct_prediction))
    optimizer.apply_gradients(zip(grads, model.trainable_variables))
    return accuracy, loss_value

  @tf.function
  def distributed_train_step(dist_inputs):
    per_replica_accuracy, per_replica_losses = strategy.run(
        train_step, args=(dist_inputs,))
    accuracy = strategy.reduce(
        tf.distribute.ReduceOp.MEAN, per_replica_accuracy, axis=None)
    loss_value = strategy.reduce(
        tf.distribute.ReduceOp.MEAN, per_replica_losses, axis=None)
    return accuracy, loss_value

  iterator = iter(ds_train)
  accuracy = 0.0
  for step in range(flags.FLAGS.train_steps):
    accuracy, loss_value = distributed_train_step(next(iterator))
    if step % display_step == 0:
      tf.print('Step %d:' % step)
      tf.print('    Loss = %f' % loss_value)
      tf.print('    Batch accuracy = %f' % accuracy)

  return accuracy
//  [python]
def main(strategy):
  """Trains an MNIST model using the given tf.distribute.Strategy."""
  # TODO(fengliuai): put this in some automatically generated code.
  os.environ[
      'TF_MLIR_TFR_LIB_DIR'] = 'tensorflow/compiler/mlir/tfr/examples/mnist'

  ds_train = tfds.load('mnist', split='train', shuffle_files=True)
  ds_train = ds_train.shuffle(1024).batch(batch_size).prefetch(64)
  ds_train = strategy.experimental_distribute_dataset(ds_train)

  with strategy.scope():
    # Create an mnist float model with the specified float state.
    model = FloatModel()
    optimizer = tf.keras.optimizers.Adam(learning_rate=learning_rate)

  def train_step(features):
    inputs = tf.image.convert_image_dtype(
        features['image'], dtype=tf.float32, saturate=False)
    labels = tf.one_hot(features['label'], num_classes)

    with tf.GradientTape() as tape:
      logits = model(inputs)
      loss_value = tf.reduce_mean(
          tf.nn.softmax_cross_entropy_with_logits(labels, logits))

    grads = tape.gradient(loss_value, model.trainable_variables)
    correct_prediction = tf.equal(tf.argmax(logits, 1), tf.argmax(labels, 1))
    accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))
    optimizer.apply_gradients(zip(grads, model.trainable_variables))
    return accuracy, loss_value

  @tf.function
  def distributed_train_step(dist_inputs):
    per_replica_accuracy, per_replica_losses = strategy.run(
        train_step, args=(dist_inputs,))
    accuracy = strategy.reduce(
        tf.distribute.ReduceOp.MEAN, per_replica_accuracy, axis=None)
    loss_value = strategy.reduce(
        tf.distribute.ReduceOp.MEAN, per_replica_losses, axis=None)
    return accuracy, loss_value

  iterator = iter(ds_train)
  accuracy = 0.0
  for step in range(flags.FLAGS.train_steps):
    accuracy, loss_value = distributed_train_step(next(iterator))
    if step % display_step == 0:
      tf.print('Step %d:' % step)
      tf.print('    Loss = %f' % loss_value)
      tf.print('    Batch accuracy = %f' % accuracy)

  return accuracy
//  [python]
def token_label_class_loss(y_true, y_pred):
    # tf.print(", y_true:", y_true.shape, "y_pred:", y_pred.shape, end="")
    if y_pred.shape[-1] != y_true.shape[-1]:
        y_pred, cls_lambda = y_pred[:, :-1], y_pred[:, -1:]
        y_true = tf.to_float(y_true)
        y_true = cls_lambda * y_true + (1 - cls_lambda) * y_true[::-1]
    return keras.losses.categorical_crossentropy(y_true, y_pred, from_logits=True)
//  [python]
def token_label_class_loss(y_true, y_pred):
    # tf.print(", y_true:", y_true.shape, "y_pred:", y_pred.shape, end="")
    if y_pred.shape[-1] != y_true.shape[-1]:
        y_pred, cls_lambda = y_pred[:, :-1], y_pred[:, -1:]
        y_true = tf.cast(y_true, y_pred.dtype)
        y_true = cls_lambda * y_true + (1 - cls_lambda) * y_true[::-1]
    return keras.losses.categorical_crossentropy(y_true, y_pred, from_logits=True)
//  [python]
def _has_fully_masked_sequence(mask):
    # Cudnn kernel will error out if the input sequence contains any
    # fully masked data. We walk around this issue by rerouting the computation
    # to standard kernel, until the issue on cudnn side has been fixed.  For a
    # fully masked sequence, it will contain all Falses. To make it easy to
    # check, we inverse the boolean, check if any of the sequence has all True.
    return tf.reduce_any(
        tf.reduce_all(tf.logical_not(t.cast(mask, dtype="bool")), axis=1)
    )
//  [python]
def _has_fully_masked_sequence(mask):
    # Cudnn kernel will error out if the input sequence contains any
    # fully masked data. We walk around this issue by rerouting the computation
    # to standard kernel, until the issue on cudnn side has been fixed.  For a
    # fully masked sequence, it will contain all Falses. To make it easy to
    # check, we inverse the boolean, check if any of the sequence has all True.
    return tf.reduce_any(
        tf.reduce_all(tf.logical_not(tf.cast(mask, dtype="bool")), axis=1)
    )
//  [python]
def losses(self):
        """
        Args:
            For `gt_classes` and `gt_anchors_deltas` parameters, see
                :meth:`RetinaNetGIou.get_ground_truth`.
            Their shapes are (N, R) and (N, R, 4), respectively, where R is
            the total number of anchors across levels, i.e. sum(Hi x Wi x A)
            For `pred_class_logits` and `pred_anchor_deltas`, see
                :meth:`RetinaNetGIouHead.forward`.

        Returns:
            dict[str: Tensor]:
                mapping from a named loss to a scalar tensor
                storing the loss. Used during training only. The dict keys are:
                "loss_cls" and "loss_box_reg"
        """
        
        assert len(self.pred_logits[0].get_shape()) == 4,"error logits dim"
        assert len(self.pred_anchor_deltas[0].get_shape()) == 4,"error anchors dim"
        
        gt_classes,gt_anchors_deltas = self._get_ground_truth()
        pred_class_logits, pred_anchor_deltas = permute_all_cls_and_box_to_N_HWA_K_and_concat(
            self.pred_logits, self.pred_anchor_deltas, self.num_classes
        )  # Shapes: (N, R, K) and (N, R, 4), respectively.


        valid_idxs = gt_classes >= 0
        foreground_idxs = (gt_classes > 0)
        num_foreground = tf.reduce_sum(t.cast(foreground_idxs,tf.int32))

        gt_classes_target = tf.boolean_mask(gt_classes,valid_idxs)
        gt_classes_target = tf.one_hot(gt_classes_target,depth=self.num_classes+1)
        gt_classes_target = gt_classes_target[:,1:]#RetinaNetGIou中没有背景, 因为背景index=0, 所以要在one hot 后去掉背景
        pred_class_logits = tf.boolean_mask(pred_class_logits,valid_idxs)

        # logits loss
        loss_cls = tf.reduce_sum(wnn.sigmoid_cross_entropy_with_logits_FL(
            labels = gt_classes_target,
            logits = pred_class_logits,
            alpha=self.focal_loss_alpha,
            gamma=self.focal_loss_gamma,
        )) / t.cast(tf.maximum(1, num_foreground),tf.float32)

        # regression loss
        pred_anchor_deltas = tf.boolean_mask(pred_anchor_deltas,foreground_idxs)
        gt_anchors_deltas = tf.boolean_mask(gt_anchors_deltas,foreground_idxs)
        B,X = wmlt.combined_static_and_dynamic_shape(foreground_idxs)
        anchors = tf.tile(self.anchors,[B,1,1])
        anchors = tf.boolean_mask(anchors,foreground_idxs)
        box = self.box2box_transform.apply_deltas(pred_anchor_deltas,anchors)
        reg_loss_sum = odl.ciou_loss(box,gt_anchors_deltas)
        loss_box_reg = tf.reduce_sum(reg_loss_sum) / t.cast(tf.maximum(1, num_foreground),tf.float32)
        loss_cls = loss_cls*self.cfg.BOX_CLS_LOSS_SCALE
        loss_box_reg = loss_box_reg*self.cfg.BOX_REG_LOSS_SCALE

        return {"loss_cls": loss_cls, "loss_box_reg": loss_box_reg}
//  [python]
def losses(self):
        """
        Args:
            For `gt_classes` and `gt_anchors_deltas` parameters, see
                :meth:`RetinaNetGIou.get_ground_truth`.
            Their shapes are (N, R) and (N, R, 4), respectively, where R is
            the total number of anchors across levels, i.e. sum(Hi x Wi x A)
            For `pred_class_logits` and `pred_anchor_deltas`, see
                :meth:`RetinaNetGIouHead.forward`.

        Returns:
            dict[str: Tensor]:
                mapping from a named loss to a scalar tensor
                storing the loss. Used during training only. The dict keys are:
                "loss_cls" and "loss_box_reg"
        """
        
        assert len(self.pred_logits[0].get_shape()) == 4,"error logits dim"
        assert len(self.pred_anchor_deltas[0].get_shape()) == 4,"error anchors dim"
        
        gt_classes,gt_anchors_deltas = self._get_ground_truth()
        pred_class_logits, pred_anchor_deltas = permute_all_cls_and_box_to_N_HWA_K_and_concat(
            self.pred_logits, self.pred_anchor_deltas, self.num_classes
        )  # Shapes: (N, R, K) and (N, R, 4), respectively.


        valid_idxs = gt_classes >= 0
        foreground_idxs = (gt_classes > 0)
        num_foreground = tf.reduce_sum(tf.cast(foreground_idxs,tf.int32))

        gt_classes_target = tf.boolean_mask(gt_classes,valid_idxs)
        gt_classes_target = tf.one_hot(gt_classes_target,depth=self.num_classes+1)
        gt_classes_target = gt_classes_target[:,1:]#RetinaNetGIou中没有背景, 因为背景index=0, 所以要在one hot 后去掉背景
        pred_class_logits = tf.boolean_mask(pred_class_logits,valid_idxs)

        # logits loss
        loss_cls = tf.reduce_sum(wnn.sigmoid_cross_entropy_with_logits_FL(
            labels = gt_classes_target,
            logits = pred_class_logits,
            alpha=self.focal_loss_alpha,
            gamma=self.focal_loss_gamma,
        )) / tf.cast(tf.maximum(1, num_foreground),tf.float32)

        # regression loss
        pred_anchor_deltas = tf.boolean_mask(pred_anchor_deltas,foreground_idxs)
        gt_anchors_deltas = tf.boolean_mask(gt_anchors_deltas,foreground_idxs)
        B,X = wmlt.combined_static_and_dynamic_shape(foreground_idxs)
        anchors = tf.tile(self.anchors,[B,1,1])
        anchors = tf.boolean_mask(anchors,foreground_idxs)
        box = self.box2box_transform.apply_deltas(pred_anchor_deltas,anchors)
        reg_loss_sum = odl.ciou_loss(box,gt_anchors_deltas)
        loss_box_reg = tf.reduce_sum(reg_loss_sum) / tf.cast(tf.maximum(1, num_foreground),tf.float32)
        loss_cls = loss_cls*self.cfg.BOX_CLS_LOSS_SCALE
        loss_box_reg = loss_box_reg*self.cfg.BOX_REG_LOSS_SCALE

        return {"loss_cls": loss_cls, "loss_box_reg": loss_box_reg}
//  [python]
def _compute_dloss_by_dmin_dmax_and_dx_symmetric(inputs: tf.Tensor,
                                                 bitwidth: tf.Tensor,
                                                 encoding_min: tf.Tensor,
                                                 encoding_max: tf.Tensor,
                                                 grad: tf.Tensor):
    """
    Calculate dloss_by_dmin, dloss_by_dmax, and dloss_by_dx tensors computed by symmetric quantization

    :param inputs: Inputs to op
    :param bitwidth: Bitwidth used to quantize
    :param encoding_min: Encoding min value(s), will be more than one if per channel is active
    :param encoding_max: Encoding max value(s), will be more than one if per channel is active
    :param grad: Gradient from child layer
    :return: Tensors for dloss_by_dmin, dloss_by_dmax, and dloss_by_dx
    """
    x = tf.to_float(inputs)
    bitwidth = tf.to_float(bitwidth)
    encoding_min = tf.to_float(encoding_min)
    encoding_max = tf.to_float(encoding_max)

    # handle min == max to avoid divide by zero
    epsilon = tf.constant(1e-5, dtype=tf.float32)
    encoding_max = tf.math.maximum(encoding_max, tf.add(encoding_min, epsilon))

    num_steps = tf.to_float(tf.pow(tf.to_float(tf.constant(2)), bitwidth) - 1)
    half_num_steps = tf.divide(num_steps, tf.constant(2.0))
    delta = encoding_max / tf.math.floor(half_num_steps)
    offset = -tf.math.ceil(half_num_steps)

    zero = tf.zeros_like(num_steps)
    x_round = tf.round(inputs / delta) - offset
    x_quant = tf.clip_by_value(x_round, zero, num_steps)

    mask_tensor = tf.to_float(tf.math.greater_equal(x_round, zero)) * \
                  tf.to_float(tf.math.less_equal(x_round, num_steps))
    grad_tensor = mask_tensor * grad

    axis = tf.cond(tf.equal(tf.rank(delta), 0),
                   lambda: tf.range(0, tf.rank(x)),         # Per-tensor
                   lambda: tf.range(0, tf.rank(x) - 1))     # Per-channel

    grad_encoding_max = tf.reduce_sum((x_quant + offset) * grad, axis=axis) - \
                        tf.reduce_sum(mask_tensor * (inputs / delta) * grad, axis=axis)

    grad_encoding_max = grad_encoding_max / tf.math.floor(half_num_steps)
    grad_encoding_max = tf.to_double(grad_encoding_max)

    return tf.negative(grad_encoding_max), grad_encoding_max, grad_tensor
//  [python]
def _compute_dloss_by_dmin_dmax_and_dx_symmetric(inputs: tf.Tensor,
                                                 bitwidth: tf.Tensor,
                                                 encoding_min: tf.Tensor,
                                                 encoding_max: tf.Tensor,
                                                 grad: tf.Tensor):
    """
    Calculate dloss_by_dmin, dloss_by_dmax, and dloss_by_dx tensors computed by symmetric quantization

    :param inputs: Inputs to op
    :param bitwidth: Bitwidth used to quantize
    :param encoding_min: Encoding min value(s), will be more than one if per channel is active
    :param encoding_max: Encoding max value(s), will be more than one if per channel is active
    :param grad: Gradient from child layer
    :return: Tensors for dloss_by_dmin, dloss_by_dmax, and dloss_by_dx
    """
    x = tf.cast(inputs, tf.float32)
    bitwidth = tf.cast(bitwidth, tf.float32)
    encoding_min = tf.cast(encoding_min, tf.float32)
    encoding_max = tf.cast(encoding_max, tf.float32)

    # handle min == max to avoid divide by zero
    epsilon = tf.constant(1e-5, dtype=tf.float32)
    encoding_max = tf.math.maximum(encoding_max, tf.add(encoding_min, epsilon))

    num_steps = tf.cast(tf.pow(tf.cast(tf.constant(2), tf.float32), bitwidth) - 1, tf.float32)
    half_num_steps = tf.divide(num_steps, tf.constant(2.0))
    delta = encoding_max / tf.math.floor(half_num_steps)
    offset = -tf.math.ceil(half_num_steps)

    zero = tf.zeros_like(num_steps)
    x_round = tf.round(inputs / delta) - offset
    x_quant = tf.clip_by_value(x_round, zero, num_steps)

    mask_tensor = tf.cast(tf.math.greater_equal(x_round, zero), tf.float32) * \
                  tf.cast(tf.math.less_equal(x_round, num_steps), tf.float32)
    grad_tensor = mask_tensor * grad

    axis = tf.cond(tf.equal(tf.rank(delta), 0),
                   lambda: tf.range(0, tf.rank(x)),         # Per-tensor
                   lambda: tf.range(0, tf.rank(x) - 1))     # Per-channel

    grad_encoding_max = tf.reduce_sum((x_quant + offset) * grad, axis=axis) - \
                        tf.reduce_sum(mask_tensor * (inputs / delta) * grad, axis=axis)

    grad_encoding_max = grad_encoding_max / tf.math.floor(half_num_steps)
    grad_encoding_max = tf.cast(grad_encoding_max, tf.float64)

    return tf.negative(grad_encoding_max), grad_encoding_max, grad_tensor
//  [python]
def _compute_dloss_by_dmin_dmax_and_dx_asymmetric(inputs: tf.Tensor,
                                                  bitwidth: tf.Tensor,
                                                  encoding_min: tf.Tensor,
                                                  encoding_max: tf.Tensor,
                                                  grad: tf.Tensor):
    """
    Calculate dloss_by_dmin, dloss_by_dmax, and dloss_by_dx tensors computed by asymmetric quantization

    :param inputs: Inputs to op
    :param bitwidth: Bitwidth used to quantize
    :param encoding_min: Encoding min value(s), will be more than one if per channel is active
    :param encoding_max: Encoding max value(s), will be more than one if per channel is active
    :param grad: Gradient from child layer
    :return: Tensors for dloss_by_dmin, dloss_by_dmax, and dloss_by_dx
    """
    x = t.cast(inputs, t.float32)
    bitwidth = t.cast(bitwidth, t.float32)
    encoding_min = t.cast(encoding_min, t.float32)
    encoding_max = t.cast(encoding_max, t.float32)

    # handle min == max to avoid divide by zero
    epsilon = t.constant(1e-5, dtype=t.float32)
    encoding_max = tf.math.maximum(encoding_max, tf.add(encoding_min, epsilon))

    num_steps = t.cast(t.pow(t.cast(t.constant(2), t.float32), bitwidth) - 1, t.float32)
    delta = (encoding_max - encoding_min) / num_steps
    b_zero = t.round(t.negative(encoding_min) / delta)
    b_zero = t.minimum(num_steps, t.maximum(t.constant(0.0), b_zero))
    offset = t.negative(b_zero)

    zero = t.zeros_like(num_steps)
    x_round = t.round(inputs / delta) - offset
    x_quant = t.clip_by_value(x_round, zero, num_steps)

    mask_tensor = t.cast(tf.math.greater_equal(x_round, zero), t.float32) * \
                  t.cast(tf.math.less_equal(x_round, num_steps), t.float32)
    grad_tensor = mask_tensor * grad

    grad_scale = (x_quant + offset - x * mask_tensor / delta) * grad
    grad_xq = delta * grad
    grad_offset = grad_xq * (1 - mask_tensor)

    axis = tf.cond(tf.equal(tf.rank(delta), 0),
                   lambda: tf.range(0, tf.rank(x)),         # Per-tensor
                   lambda: tf.range(0, tf.rank(x) - 1))     # Per-channel

    intermediate_term1 = tf.reduce_sum(grad_scale, axis=axis) / num_steps
    intermediate_term2 = num_steps / (encoding_max - encoding_min) ** 2 * tf.reduce_sum(grad_offset, axis=axis)

    grad_encoding_min = -intermediate_term1 + encoding_max * intermediate_term2
    grad_encoding_max = intermediate_term1 - encoding_min * intermediate_term2

    grad_encoding_max = t.cast(grad_encoding_max, t.float64)
    grad_encoding_min = t.cast(grad_encoding_min, t.float64)

    return grad_encoding_min, grad_encoding_max, grad_tensor
//  [python]
def _compute_dloss_by_dmin_dmax_and_dx_asymmetric(inputs: tf.Tensor,
                                                  bitwidth: tf.Tensor,
                                                  encoding_min: tf.Tensor,
                                                  encoding_max: tf.Tensor,
                                                  grad: tf.Tensor):
    """
    Calculate dloss_by_dmin, dloss_by_dmax, and dloss_by_dx tensors computed by asymmetric quantization

    :param inputs: Inputs to op
    :param bitwidth: Bitwidth used to quantize
    :param encoding_min: Encoding min value(s), will be more than one if per channel is active
    :param encoding_max: Encoding max value(s), will be more than one if per channel is active
    :param grad: Gradient from child layer
    :return: Tensors for dloss_by_dmin, dloss_by_dmax, and dloss_by_dx
    """
    x = tf.cast(inputs, tf.float32)
    bitwidth = tf.cast(bitwidth, tf.float32)
    encoding_min = tf.cast(encoding_min, tf.float32)
    encoding_max = tf.cast(encoding_max, tf.float32)

    # handle min == max to avoid divide by zero
    epsilon = tf.constant(1e-5, dtype=tf.float32)
    encoding_max = tf.math.maximum(encoding_max, tf.add(encoding_min, epsilon))

    num_steps = tf.cast(tf.pow(tf.cast(tf.constant(2), tf.float32), bitwidth) - 1, tf.float32)
    delta = (encoding_max - encoding_min) / num_steps
    b_zero = tf.round(tf.negative(encoding_min) / delta)
    b_zero = tf.minimum(num_steps, tf.maximum(tf.constant(0.0), b_zero))
    offset = tf.negative(b_zero)

    zero = tf.zeros_like(num_steps)
    x_round = tf.round(inputs / delta) - offset
    x_quant = tf.clip_by_value(x_round, zero, num_steps)

    mask_tensor = tf.cast(tf.math.greater_equal(x_round, zero), tf.float32) * \
                  tf.cast(tf.math.less_equal(x_round, num_steps), tf.float32)
    grad_tensor = mask_tensor * grad

    grad_scale = (x_quant + offset - x * mask_tensor / delta) * grad
    grad_xq = delta * grad
    grad_offset = grad_xq * (1 - mask_tensor)

    axis = tf.cond(tf.equal(tf.rank(delta), 0),
                   lambda: tf.range(0, tf.rank(x)),         # Per-tensor
                   lambda: tf.range(0, tf.rank(x) - 1))     # Per-channel

    intermediate_term1 = tf.reduce_sum(grad_scale, axis=axis) / num_steps
    intermediate_term2 = num_steps / (encoding_max - encoding_min) ** 2 * tf.reduce_sum(grad_offset, axis=axis)

    grad_encoding_min = -intermediate_term1 + encoding_max * intermediate_term2
    grad_encoding_max = intermediate_term1 - encoding_min * intermediate_term2

    grad_encoding_max = tf.cast(grad_encoding_max, tf.float64)
    grad_encoding_min = tf.cast(grad_encoding_min, tf.float64)

    return grad_encoding_min, grad_encoding_max, grad_tensor
//  [python]
def _is_sequence_right_padded(mask):
    """Check the mask tensor and see if it right padded.

    cuDNN uses the sequence length param to skip the tailing
    timestep. If the data is left padded, or not a strict right padding (has
    masked value in the middle of the sequence), then cuDNN won't work
    properly in those cases.

    Left padded data: [[False, False, True, True, True]].
    Right padded data: [[True, True, True, False, False]].
    Mixture of mask/unmasked data: [[True, False, True, False, False]].

    Note that for the mixed data example above, the actually data RNN should see
    are those 2 Trues (index 0 and 2), the index 1 False should be ignored and
    not pollute the internal states.

    Args:
        mask: the Boolean tensor with shape [batch, timestep]

    Returns:
        boolean scalar tensor, whether the mask is strictly right padded.
    """
    max_seq_length = tf.shape(mask)[1]
    count_of_true = tf.reduce_sum(tf.to_int32(mask), axis=1)
    right_padded_mask = tf.sequence_mask(count_of_true, maxlen=max_seq_length)
    return tf.reduce_all(
        tf.equal(
            tf.to_bool(mask),
            tf.to_bool(right_padded_mask),
        )
    )