jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def _sample_discrete_gaussian_helper(scale, shape, dtype):
  """Draw samples from discrete Gaussian, assuming scale >= 0."""
  scale = tf.cast(scale, tf.int64)
  sq_scale = _int_square(scale)

  # Scale for discrete Laplace. The sampling algorithm should be correct
  # for any discrete Laplace scale, and the original paper uses
  # `dlap_scale = floor(scale) + 1`. Here we use `dlap_scale = scale` (where
  # input `scale` is restricted to integers >= 1) to simplify the fraction
  # below. It turns out that for integer scales >= 1, `dlap_scale = scale` gives
  # a good minimum success rate of ~70%, allowing a small oversampling factor.
  dlap_scale = scale
  oversample_factor = 1.5

  # Draw at least some samples in case we got unlucky with small input shape.
  min_n = 1000
  target_n = tf.reduce_prod(tf.cast(shape, tf.int64))
  oversample_n = oversample_factor * tf.cast(target_n, tf.float32)
  draw_n = tf.maximum(min_n, tf.cast(oversample_n, tf.int32))

  accepted_n = tf.constant(0, dtype=target_n.dtype)
  result = tf.zeros((0,), dtype=tf.int64)

  while accepted_n < target_n:
    # Since the number of samples could be different in every retry, we need to
    # manually specify the shape info for TF.
    tf.autograph.experimental.set_loop_options(
        shape_invariants=[(result, tf.TensorShape([None]))])
    # Draw samples.
    samples = _sample_discrete_laplace(dlap_scale, shape=(draw_n,))
    z_numer = _int_square(tf.abs(samples) - scale)
    z_denom = 2 * sq_scale
    bern_probs = tf.exp(-1.0 * tf.divide(z_numer, z_denom))
    accept = _sample_bernoulli(bern_probs)
    # Keep successful samples and increment counter.
    accepted_samples = samples[tf.equal(accept, 1)]
    accepted_n += tf.cast(tf.size(accepted_samples), accepted_n.dtype)
    result = tf.concat([result, accepted_samples], axis=0)
    # Reduce the number of draws for any retries.
    draw_n = tf.cast(target_n - accepted_n, tf.float32) * oversample_factor
    draw_n = tf.maximum(min_n, tf.cast(draw_n, tf.int32))

  return tf.cast(tf.reshape(result[:target_n], shape), dtype)
//  [python]
def normalized_to_image_coordinates(normalized_boxes, image_shape,
                                    parallel_iterations=32):
  """Converts a batch of boxes from normal to image coordinates.

  Args:
    normalized_boxes: a float32 tensor of shape [None, num_boxes, 4] in
      normalized coordinates.
    image_shape: a float32 tensor of shape [4] containing the image shape.
    parallel_iterations: parallelism for the map_fn op.

  Returns:
    absolute_boxes: a float32 tensor of shape [None, num_boxes, 4] containing
      the boxes in image coordinates.
  """
  x_scale = t.cast(image_shape[2], t.float32)
  y_scale = t.cast(image_shape[1], t.float32)
  def _to_absolute_coordinates(normalized_boxes):
    y_min, x_min, y_max, x_max = tf.split(
        value=normalized_boxes, num_or_size_splits=4, axis=1)
    y_min = y_scale * y_min
    y_max = y_scale * y_max
    x_min = x_scale * x_min
    x_max = x_scale * x_max
    scaled_boxes = tf.concat([y_min, x_min, y_max, x_max], 1)
    return scaled_boxes

  absolute_boxes = shape_utils.static_or_dynamic_map_fn(
      _to_absolute_coordinates,
      elems=(normalized_boxes),
      dtype=tf.float32,
      parallel_iterations=parallel_iterations,
      back_prop=True)
  return absolute_boxes
//  [python]
def normalized_to_image_coordinates(normalized_boxes, image_shape,
                                    parallel_iterations=32):
  """Converts a batch of boxes from normal to image coordinates.

  Args:
    normalized_boxes: a float32 tensor of shape [None, num_boxes, 4] in
      normalized coordinates.
    image_shape: a float32 tensor of shape [4] containing the image shape.
    parallel_iterations: parallelism for the map_fn op.

  Returns:
    absolute_boxes: a float32 tensor of shape [None, num_boxes, 4] containing
      the boxes in image coordinates.
  """
  x_scale = tf.cast(image_shape[2], tf.float32)
  y_scale = tf.cast(image_shape[1], tf.float32)
  def _to_absolute_coordinates(normalized_boxes):
    y_min, x_min, y_max, x_max = tf.split(
        value=normalized_boxes, num_or_size_splits=4, axis=1)
    y_min = y_scale * y_min
    y_max = y_scale * y_max
    x_min = x_scale * x_min
    x_max = x_scale * x_max
    scaled_boxes = tf.concat([y_min, x_min, y_max, x_max], 1)
    return scaled_boxes

  absolute_boxes = shape_utils.static_or_dynamic_map_fn(
      _to_absolute_coordinates,
      elems=(normalized_boxes),
      dtype=tf.float32,
      parallel_iterations=parallel_iterations,
      back_prop=True)
  return absolute_boxes
//  [python]
def test_orphan_nodes(self):
    """Tests sampling subgraph from nodes without edges."""
    (unused_eats_edges, sampler, graph_schema,
     sampling_spec) = _get_test_edges_sampler_schema_spec()
    pipeline = subgraph_pipeline.SamplingPipeline(
        graph_schema, sampling_spec, sampler,
        node_features_accessor_factory=_test_create_feature_accessor)

    seed_nodes = tf.ragged.constant([
        ['unicorn'],  # Orphan node (no edges!)
        ['unicorn', 'dog'],  # Orphan node with a non-orphan node.
    ])
    graph_tensor = pipeline(seed_nodes)
    # Break into its subgraphs.
    graph_tensors = list(tf.data.Dataset.from_tensors(graph_tensor).unbatch())
    graph0 = graph_tensors[0]
    graph1 = graph_tensors[1]

    self.assertAllEqual(graph0.node_sets['animal']['#id'], [b'unicorn'])
    # No edges (unicorn needs no food).
    for edge_set0 in graph0.edge_sets.values():
      self.assertAllEqual(edge_set0.sizes, [0])

    # Dog eats 2 things (unicorn eats none):
    self.assertAllEqual(graph1.edge_sets['eats'].sizes, [2])

    # Make sure 'unicorn' appears in node set but without edges.
    bool_mask = graph1.node_sets['animal']['#id'] == b'unicorn'
    self.assertAllEqual(tf.reduce_sum(tf.cast(bool_mask, tf.int32)), 1)
    unicorn_index = tf.argmax(bool_mask, output_type=tf.int32)
    self.assertNotIn(unicorn_index, graph1.edge_sets['eats'].adjacency.source)
//  [python]
def update_state(self, y_true, y_pred, sample_weight=None):
    """Updates `m`, `nrow`, `ncol` respectively the joint and
        marginal histograms of (`y_true`, `y_pred`)
        """

    y_true = tf.clip_by_value(y_true, self.actual_min, self.actual_max)
    y_pred = tf.clip_by_value(y_pred, self.preds_min, self.preds_max)

    i = (
        tf.searchsorted(
            self.actual_cuts,
            tf.reshape(y_true, [-1]),
            side="right",
            out_type=tf.int64,
        ) - 1
    )
    j = (
        tf.searchsorted(
            self.preds_cuts,
            tf.reshape(y_pred, [-1]),
            side="right",
            out_type=tf.int64,
        ) - 1
    )

    nrow = tf.tensor_scatter_nd_add(self.nrow, tf.expand_dims(i, axis=-1), tf.ones_like(i))
    ncol = tf.tensor_scatter_nd_add(self.ncol, tf.expand_dims(j, axis=-1), tf.ones_like(j))
    ij = tf.stack([i, j], axis=1)
    m = tf.tensor_scatter_nd_add(self.m, ij, tf.ones_like(i))

    self.n.assign_add(tf.shape(i, out_type=tf.int64)[0])
    self.m.assign(m)
    self.nrow.assign(nrow)
    self.ncol.assign(ncol)
//  [python]
def update_state(self, y_true, y_pred, sample_weight=None):
    """Updates `m`, `nrow`, `ncol` respectively the joint and
        marginal histograms of (`y_true`, `y_pred`)
        """

    y_true = tf.clip_by_value(y_true, self.actual_min, self.actual_max)
    y_pred = tf.clip_by_value(y_pred, self.preds_min, self.preds_max)

    i = (
        tf.searchsorted(
            self.actual_cuts,
            tf.cast(tf.reshape(y_true, [-1]), self.actual_cuts.dtype),
            side="right",
            out_type=tf.int64,
        ) - 1
    )
    j = (
        tf.searchsorted(
            self.preds_cuts,
            tf.cast(tf.reshape(y_pred, [-1]), self.preds_cuts.dtype),
            side="right",
            out_type=tf.int64,
        ) - 1
    )

    nrow = tf.tensor_scatter_nd_add(self.nrow, tf.expand_dims(i, axis=-1), tf.ones_like(i))
    ncol = tf.tensor_scatter_nd_add(self.ncol, tf.expand_dims(j, axis=-1), tf.ones_like(j))
    ij = tf.stack([i, j], axis=1)
    m = tf.tensor_scatter_nd_add(self.m, ij, tf.ones_like(i))

    self.n.assign_add(tf.shape(i, out_type=tf.int64)[0])
    self.m.assign(m)
    self.nrow.assign(nrow)
    self.ncol.assign(ncol)
//  [python]
def result(self):
    p, q, _, _ = self._compute_variables()
    n = t.cast(self.n, t.float32)
    non_zeros_col = t.math.count_nonzero(self.ncol)
    non_zeros_row = t.math.count_nonzero(self.nrow)
    m = t.cast(t.minimum(non_zeros_col, non_zeros_row), t.float32)
    return 2 * (p - q) / (t.square(n) * (m - 1) / m)
//  [python]
def result(self):
    p, q, _, _ = self._compute_variables()
    n = tf.cast(self.n, tf.float32)
    non_zeros_col = tf.math.count_nonzero(self.ncol)
    non_zeros_row = tf.math.count_nonzero(self.nrow)
    m = tf.cast(tf.minimum(non_zeros_col, non_zeros_row), tf.float32)
    return 2 * (p - q) / (tf.square(n) * (m - 1) / m)
//  [python]
def _pad_to_window_size(
        self,
        input_ids,
        attention_mask,
        token_type_ids,
        position_ids,
        inputs_embeds,
        pad_token_id,
    ):
        """A helper function to pad tokens and mask to work with implementation of Longformer selfattention."""
        # padding
        attention_window = (
            self.attention_window if isinstance(self.attention_window, int) else max(self.attention_window)
        )

        assert attention_window % 2 == 0, f"`attention_window` should be an even value. Given {attention_window}"

        input_shape = shape_list(input_ids) if input_ids is not None else shape_list(inputs_embeds)
        batch_size, seq_len = input_shape[:2]
        padding_len = (attention_window - seq_len % attention_window) % attention_window

        if padding_len > 0:
            logger.info(
                f"Input ids are automatically padded from {seq_len} to {seq_len + padding_len} to be a multiple of "
                f"`config.attention_window`: {attention_window}"
            )

        paddings = tf.convert_to_tensor([[0, 0], [0, padding_len]])

        if input_ids is not None:
            input_ids = tf.pad(input_ids, paddings, constant_values=pad_token_id)

        if position_ids is not None:
            # pad with position_id = pad_token_id as in modeling_roberta.RobertaEmbeddings
            position_ids = tf.pad(position_ids, paddings, constant_values=pad_token_id)

        if inputs_embeds is not None:
            if padding_len > 0:
                input_ids_padding = tf.fill((batch_size, padding_len), self.pad_token_id)
                inputs_embeds_padding = self.embeddings(input_ids_padding)
                inputs_embeds = tf.concat([inputs_embeds, inputs_embeds_padding], axis=-2)

        attention_mask = tf.pad(attention_mask, paddings, constant_values=False)  # no attention on the padding tokens
        token_type_ids = tf.pad(token_type_ids, paddings, constant_values=0)  # pad with token_type_id = 0

        return (
            padding_len,
            input_ids,
            attention_mask,
            token_type_ids,
            position_ids,
            inputs_embeds,
        )
//  [python]
def _pad_to_window_size(
        self,
        input_ids,
        attention_mask,
        token_type_ids,
        position_ids,
        inputs_embeds,
        pad_token_id,
    ):
        """A helper function to pad tokens and mask to work with implementation of Longformer selfattention."""
        # padding
        attention_window = (
            self.attention_window if isinstance(self.attention_window, int) else max(self.attention_window)
        )

        assert attention_window % 2 == 0, f"`attention_window` should be an even value. Given {attention_window}"

        input_shape = shape_list(input_ids) if input_ids is not None else shape_list(inputs_embeds)
        batch_size, seq_len = input_shape[:2]
        padding_len = (attention_window - seq_len % attention_window) % attention_window

        if padding_len > 0:
            logger.info(
                f"Input ids are automatically padded from {seq_len} to {seq_len + padding_len} to be a multiple of "
                f"`config.attention_window`: {attention_window}"
            )

        paddings = tf.convert_to_tensor([[0, 0], [0, padding_len]])

        if input_ids is not None:
            input_ids = tf.pad(input_ids, paddings, constant_values=pad_token_id)

        if position_ids is not None:
            # pad with position_id = pad_token_id as in modeling_roberta.RobertaEmbeddings
            position_ids = tf.pad(position_ids, paddings, constant_values=pad_token_id)

        if inputs_embeds is not None:
            if padding_len > 0:
                input_ids_padding = tf.cast(tf.fill((batch_size, padding_len), self.pad_token_id), tf.int64)
                inputs_embeds_padding = self.embeddings(input_ids_padding)
                inputs_embeds = tf.concat([inputs_embeds, inputs_embeds_padding], axis=-2)

        attention_mask = tf.pad(attention_mask, paddings, constant_values=False)  # no attention on the padding tokens
        token_type_ids = tf.pad(token_type_ids, paddings, constant_values=0)  # pad with token_type_id = 0

        return (
            padding_len,
            input_ids,
            attention_mask,
            token_type_ids,
            position_ids,
            inputs_embeds,
        )
//  [python]
def _compute_variables(self):
    """Compute a tuple containing the concordant pairs, discordant pairs,
        ties in `y_true` and `y_pred`.

        Returns:
          A tuple
        """
    m = t.cast(self.m, t.float32)
    n_cap = tf.cumsum(tf.cumsum(m, axis=0), axis=1)
    # Number of concordant pairs.
    p = tf.math.reduce_sum(tf.multiply(n_cap[:-1, :-1], m[1:, 1:]))
    sum_m_squard = tf.math.reduce_sum(tf.math.square(m))
    # Ties in x.
    t = (t.cast(tf.math.reduce_sum(tf.math.square(self.nrow)), t.float32) - sum_m_squard) / 2.0
    # Ties in y.
    u = (t.cast(tf.math.reduce_sum(tf.math.square(self.ncol)), t.float32) - sum_m_squard) / 2.0
    # Ties in both.
    b = tf.math.reduce_sum(tf.multiply(m, (m - 1.0))) / 2.0
    # Number of discordant pairs.
    n = t.cast(self.n, t.float32)
    q = (n - 1.0) * n / 2.0 - p - t - u - b
    return p, q, t, u
//  [python]
def _compute_variables(self):
    """Compute a tuple containing the concordant pairs, discordant pairs,
        ties in `y_true` and `y_pred`.

        Returns:
          A tuple
        """
    m = tf.cast(self.m, tf.float32)
    n_cap = tf.cumsum(tf.cumsum(m, axis=0), axis=1)
    # Number of concordant pairs.
    p = tf.math.reduce_sum(tf.multiply(n_cap[:-1, :-1], m[1:, 1:]))
    sum_m_squard = tf.math.reduce_sum(tf.math.square(m))
    # Ties in x.
    t = (tf.cast(tf.math.reduce_sum(tf.math.square(self.nrow)), tf.float32) - sum_m_squard) / 2.0
    # Ties in y.
    u = (tf.cast(tf.math.reduce_sum(tf.math.square(self.ncol)), tf.float32) - sum_m_squard) / 2.0
    # Ties in both.
    b = tf.math.reduce_sum(tf.multiply(m, (m - 1.0))) / 2.0
    # Number of discordant pairs.
    n = tf.cast(self.n, tf.float32)
    q = (n - 1.0) * n / 2.0 - p - t - u - b
    return p, q, t, u
//  [python]
def _apply_dense(self, grad, var):
    m = self.get_slot(var, "m")
    v = self.get_slot(var, "v")

    lr = t.cast(self._lr_t, grad.dtype.base_dtype)
    beta1 = t.cast(self._beta1_t, grad.dtype.base_dtype)
    beta2 = t.cast(self._beta2_t, grad.dtype.base_dtype)
    epsilon = t.cast(self._epsilon_t, grad.dtype.base_dtype)

    grad = grad - var * self._l2_weight_decay

    # m_t = beta_1 * m_{t-1} + (1-beta_1) * g_t
    m_t = m.assign(beta1 * m + (1.0 - beta1) * grad)

    # v_t = beta_2 * v_{t-1} + (1-beta_2) * g_t ** 2
    v_t = v.assign(beta2 * v + (1.0 - beta2) * grad * grad)

    if self._use_bias_correction:
      beta1_power, beta2_power = self._get_beta_accumulators()
      beta1_power = t.cast(beta1_power, grad.dtype.base_dtype)
      beta2_power = t.cast(beta2_power, grad.dtype.base_dtype)
      lr_t = lr * tf.sqrt(1.0 - beta2_power) / (1.0 - beta1_power)
    else:
      lr_t = lr

    if self._use_nesterov:
      # delta theta = lr_t * (
      #    (beta_1 * m_t + (1-beta1) * g_t) / (sqrt(v_t) + epsilon))
      step = lr_t * ((beta1 * m_t + (1.0 - beta1) * grad) /
                     (tf.sqrt(v_t) + epsilon))
    else:
      # delta theta = lr_t * m_t / (sqrt(v_t) + epsilon)
      step = lr_t * m_t / (tf.sqrt(v_t) + epsilon)

    # AdamW style weight decay term.
    step = step + lr_t * self._adamw_weight_decay * var

    theta_t = tf1.assign_sub(var, step)

    return tf.group(*[theta_t, m_t, v_t])
//  [python]
def _apply_dense(self, grad, var):
    m = self.get_slot(var, "m")
    v = self.get_slot(var, "v")

    lr = tf.cast(self._lr_t, grad.dtype.base_dtype)
    beta1 = tf.cast(self._beta1_t, grad.dtype.base_dtype)
    beta2 = tf.cast(self._beta2_t, grad.dtype.base_dtype)
    epsilon = tf.cast(self._epsilon_t, grad.dtype.base_dtype)

    grad = grad - var * self._l2_weight_decay

    # m_t = beta_1 * m_{t-1} + (1-beta_1) * g_t
    m_t = m.assign(beta1 * m + (1.0 - beta1) * grad)

    # v_t = beta_2 * v_{t-1} + (1-beta_2) * g_t ** 2
    v_t = v.assign(beta2 * v + (1.0 - beta2) * grad * grad)

    if self._use_bias_correction:
      beta1_power, beta2_power = self._get_beta_accumulators()
      beta1_power = tf.cast(beta1_power, grad.dtype.base_dtype)
      beta2_power = tf.cast(beta2_power, grad.dtype.base_dtype)
      lr_t = lr * tf.sqrt(1.0 - beta2_power) / (1.0 - beta1_power)
    else:
      lr_t = lr

    if self._use_nesterov:
      # delta theta = lr_t * (
      #    (beta_1 * m_t + (1-beta1) * g_t) / (sqrt(v_t) + epsilon))
      step = lr_t * ((beta1 * m_t + (1.0 - beta1) * grad) /
                     (tf.sqrt(v_t) + epsilon))
    else:
      # delta theta = lr_t * m_t / (sqrt(v_t) + epsilon)
      step = lr_t * m_t / (tf.sqrt(v_t) + epsilon)

    # AdamW style weight decay term.
    step = step + lr_t * self._adamw_weight_decay * var

    theta_t = tf1.assign_sub(var, step)

    return tf.group(*[theta_t, m_t, v_t])
//  [python]
def result(self):
    nrow = t.cast(self.nrow, tf.float32)
    ncol = t.cast(self.ncol, tf.float32)
    n = t.cast(self.n, tf.float32)

    nrow_ = tf.where(nrow > 0, nrow, -1.0)
    rrow = tf.cumsum(nrow, exclusive=True) + (nrow_ - n) / 2
    ncol_ = tf.where(ncol > 0, ncol, -1.0)
    rcol = tf.cumsum(ncol, exclusive=True) + (ncol_ - n) / 2

    rrow = rrow / tf.math.sqrt(tf.reduce_sum(nrow * tf.square(rrow)))
    rcol = rcol / tf.math.sqrt(tf.reduce_sum(ncol * tf.square(rcol)))

    m = t.cast(self.m, tf.float32)
    corr = tf.matmul(tf.expand_dims(rrow, axis=0), m)
    corr = tf.matmul(corr, tf.expand_dims(rcol, axis=1))
    return tf.squeeze(corr)
//  [python]
def result(self):
    nrow = tf.cast(self.nrow, tf.float32)
    ncol = tf.cast(self.ncol, tf.float32)
    n = tf.cast(self.n, tf.float32)

    nrow_ = tf.where(nrow > 0, nrow, -1.0)
    rrow = tf.cumsum(nrow, exclusive=True) + (nrow_ - n) / 2
    ncol_ = tf.where(ncol > 0, ncol, -1.0)
    rcol = tf.cumsum(ncol, exclusive=True) + (ncol_ - n) / 2

    rrow = rrow / tf.math.sqrt(tf.reduce_sum(nrow * tf.square(rrow)))
    rcol = rcol / tf.math.sqrt(tf.reduce_sum(ncol * tf.square(rcol)))

    m = tf.cast(self.m, tf.float32)
    corr = tf.matmul(tf.expand_dims(rrow, axis=0), m)
    corr = tf.matmul(corr, tf.expand_dims(rcol, axis=1))
    return tf.squeeze(corr)
//  [python]
def __init__(
      self,
      actual_min: float = 0.0,
      actual_max: float = 1.0,
      preds_min: float = 0.0,
      preds_max: float = 1.0,
      actual_cutpoints: int = 100,
      preds_cutpoints: int = 100,
      name: str = None,
      dtype: AcceptableDTypes = None,
  ):
    """Creates a `CorrelationBase` instance.

        Args:
          actual_min: the inclusive lower bound on values from actual.
          actual_max: the exclusive upper bound on values from actual.
          preds_min: the inclusive lower bound on values from preds.
          preds_max: the exclusive upper bound on values from preds.
          actual_cutpoints: the number of divisions to create in actual range,
            defaults to 100.
          preds_cutpoints: the number of divisions to create in preds range,
            defaults to 100.
          name: (optional) String name of the metric instance
          dtype: (optional) Data type of the metric result. Defaults to `None`
        """
    super().__init__(name=name, dtype=dtype)
    self.actual_min = actual_min
    self.actual_max = actual_max
    self.preds_min = preds_min
    self.preds_max = preds_max
    self.actual_cutpoints = actual_cutpoints
    self.preds_cutpoints = preds_cutpoints
    actual_cuts = np.linspace(
        t.cast(self.actual_min, t.float32),
        t.cast(self.actual_max, t.float32),
        self.actual_cutpoints,
    )
    actual_cuts[-1] += backend.epsilon()
    preds_cuts = np.linspace(
        t.cast(self.preds_min, t.float32),
        t.cast(self.preds_max, t.float32),
        self.preds_cutpoints,
    )
    preds_cuts[-1] += backend.epsilon()
    self.actual_cuts = tf.convert_to_tensor(actual_cuts, tf.float32)
    self.preds_cuts = tf.convert_to_tensor(preds_cuts, tf.float32)
    self.m = self.add_weight("m", (self.actual_cutpoints - 1, self.preds_cutpoints - 1), dtype=tf.int64)
    self.nrow = self.add_weight("nrow", (self.actual_cutpoints - 1), dtype=tf.int64)
    self.ncol = self.add_weight("ncol", (self.preds_cutpoints - 1), dtype=tf.int64)
    self.n = self.add_weight("n", (), dtype=tf.int64)
//  [python]
def __init__(
      self,
      actual_min: float = 0.0,
      actual_max: float = 1.0,
      preds_min: float = 0.0,
      preds_max: float = 1.0,
      actual_cutpoints: int = 100,
      preds_cutpoints: int = 100,
      name: str = None,
      dtype: AcceptableDTypes = None,
  ):
    """Creates a `CorrelationBase` instance.

        Args:
          actual_min: the inclusive lower bound on values from actual.
          actual_max: the exclusive upper bound on values from actual.
          preds_min: the inclusive lower bound on values from preds.
          preds_max: the exclusive upper bound on values from preds.
          actual_cutpoints: the number of divisions to create in actual range,
            defaults to 100.
          preds_cutpoints: the number of divisions to create in preds range,
            defaults to 100.
          name: (optional) String name of the metric instance
          dtype: (optional) Data type of the metric result. Defaults to `None`
        """
    super().__init__(name=name, dtype=dtype)
    self.actual_min = actual_min
    self.actual_max = actual_max
    self.preds_min = preds_min
    self.preds_max = preds_max
    self.actual_cutpoints = actual_cutpoints
    self.preds_cutpoints = preds_cutpoints
    actual_cuts = np.linspace(
        tf.cast(self.actual_min, tf.float32),
        tf.cast(self.actual_max, tf.float32),
        self.actual_cutpoints,
    )
    actual_cuts[-1] += backend.epsilon()
    preds_cuts = np.linspace(
        tf.cast(self.preds_min, tf.float32),
        tf.cast(self.preds_max, tf.float32),
        self.preds_cutpoints,
    )
    preds_cuts[-1] += backend.epsilon()
    self.actual_cuts = tf.convert_to_tensor(actual_cuts, tf.float32)
    self.preds_cuts = tf.convert_to_tensor(preds_cuts, tf.float32)
    self.m = self.add_weight("m", (self.actual_cutpoints - 1, self.preds_cutpoints - 1), dtype=tf.int64)
    self.nrow = self.add_weight("nrow", (self.actual_cutpoints - 1), dtype=tf.int64)
    self.ncol = self.add_weight("ncol", (self.preds_cutpoints - 1), dtype=tf.int64)
    self.n = self.add_weight("n", (), dtype=tf.int64)
//  [python]
def _maybe_cast(tensor, dtype):
    if tensor.dtype != dtype:
      return t.cast(tensor, dtype)
    else:
      return tensor
//  [python]
def _maybe_cast(tensor, dtype):
    if tensor.dtype != dtype:
      return tf.cast(tensor, dtype)
    else:
      return tensor
//  [python]
def result(self):
    ncol = t.cast(self.ncol, t.float32)
    nrow = t.cast(self.nrow, t.float32)
    m = t.cast(self.m, t.float32)
    n = t.cast(self.n, t.float32)

    col_bins = (self.preds_cuts[1:] - self.preds_cuts[:-1]) / 2.0 + self.preds_cuts[:-1]
    row_bins = (self.actual_cuts[1:] - self.actual_cuts[:-1]) / 2.0 + self.actual_cuts[:-1]

    n_col = t.reduce_sum(ncol)
    n_row = t.reduce_sum(nrow)
    col_mean = t.reduce_sum(ncol * col_bins) / n_col
    row_mean = t.reduce_sum(nrow * row_bins) / n_row

    col_var = t.reduce_sum(ncol * t.square(col_bins)) - n_col * t.square(col_mean)
    row_var = t.reduce_sum(nrow * t.square(row_bins)) - n_row * t.square(row_mean)

    joint_product = m * t.expand_dims(row_bins, axis=1) * col_bins

    corr = t.reduce_sum(joint_product) - n * col_mean * row_mean

    return corr / t.sqrt(col_var * row_var)
//  [python]
def result(self):
    ncol = tf.cast(self.ncol, tf.float32)
    nrow = tf.cast(self.nrow, tf.float32)
    m = tf.cast(self.m, tf.float32)
    n = tf.cast(self.n, tf.float32)

    col_bins = (self.preds_cuts[1:] - self.preds_cuts[:-1]) / 2.0 + self.preds_cuts[:-1]
    row_bins = (self.actual_cuts[1:] - self.actual_cuts[:-1]) / 2.0 + self.actual_cuts[:-1]

    n_col = tf.reduce_sum(ncol)
    n_row = tf.reduce_sum(nrow)
    col_mean = tf.reduce_sum(ncol * col_bins) / n_col
    row_mean = tf.reduce_sum(nrow * row_bins) / n_row

    col_var = tf.reduce_sum(ncol * tf.square(col_bins)) - n_col * tf.square(col_mean)
    row_var = tf.reduce_sum(nrow * tf.square(row_bins)) - n_row * tf.square(row_mean)

    joint_product = m * tf.expand_dims(row_bins, axis=1) * col_bins

    corr = tf.reduce_sum(joint_product) - n * col_mean * row_mean

    return corr / tf.sqrt(col_var * row_var)
//  [python]
def _safe_one_hot_encoding(tensor, num_classes, dtype=None):
    """Given a (possibly out of range) vector of labels, transform them into one-hot encoding."""
    one_hot_encoded = slim.one_hot_encoding(
      tensor, num_classes, on_value=tf.constant(1, tf.int64),
      off_value=tf.constant(0, tf.int64))

    # This makes sure that when there are no labels, reduce_max will not output -inf but output 0s instead.
    stacked = tf.concat((tf.zeros((1, num_classes), dtype=one_hot_encoded.dtype),
                         one_hot_encoded), axis=0)
    tensor = tf.reduce_max(stacked, 0)
    if dtype is not None:
      tensor = tf.to_int32(tensor)  # GeneralModel._dtype_string_to_dtype(FLAGS.dataset_dtype)
    return tensor
//  [python]
def _safe_one_hot_encoding(tensor, num_classes, dtype=None):
    """Given a (possibly out of range) vector of labels, transform them into one-hot encoding."""
    one_hot_encoded = slim.one_hot_encoding(
      tensor, num_classes, on_value=tf.constant(1, tf.int64),
      off_value=tf.constant(0, tf.int64))

    # This makes sure that when there are no labels, reduce_max will not output -inf but output 0s instead.
    stacked = tf.concat((tf.zeros((1, num_classes), dtype=one_hot_encoded.dtype),
                         one_hot_encoded), axis=0)
    tensor = tf.reduce_max(stacked, 0)
    if dtype is not None:
      tensor = tf.cast(tensor, dtype)  # GeneralModel._dtype_string_to_dtype(FLAGS.dataset_dtype)
    return tensor
//  [python]
def next_inputs(
        self,
        time,
        outputs,
        state,
        sample_ids,
        stop_token_prediction,
        training=False,
        **kwargs,
    ):
        if training:
            finished = time + 1 >= self.max_lengths
            next_inputs = (
                self._ratio * self.targets[:, time, :]
                + (1.0 - self._ratio) * outputs[:, -self.config.n_mels :]
            )
            next_state = state
            return (finished, next_inputs, next_state)
        else:
            stop_token_prediction = tf.nn.sigmoid(stop_token_prediction)
            finished = tf.round(stop_token_prediction)
            finished = tf.reduce_all(finished)
            next_inputs = outputs[:, -self.config.n_mels :]
            next_state = state
            return (finished, next_inputs, next_state)
//  [python]
def next_inputs(
        self,
        time,
        outputs,
        state,
        sample_ids,
        stop_token_prediction,
        training=False,
        **kwargs,
    ):
        if training:
            finished = time + 1 >= self.max_lengths
            next_inputs = (
                self._ratio * self.targets[:, time, :]
                + (1.0 - self._ratio) * outputs[:, -self.config.n_mels :]
            )
            next_state = state
            return (finished, next_inputs, next_state)
        else:
            stop_token_prediction = tf.nn.sigmoid(stop_token_prediction)
            finished = tf.cast(tf.round(stop_token_prediction), tf.bool)
            finished = tf.reduce_all(finished)
            next_inputs = outputs[:, -self.config.n_mels :]
            next_state = state
            return (finished, next_inputs, next_state)
//  [python]
def _compute_window_mask(self, max_alignments):
        """Compute window mask for inference.
        Args:
            max_alignments (int): [batch_size]
        """
        expanded_max_alignments = tf.expand_dims(max_alignments, 1)  # [batch_size, 1]
        low = expanded_max_alignments - self.win_front
        high = expanded_max_alignments + self.win_back
        mlow = tf.to_float((self._indices < low))
        mhigh = tf.to_float((self._indices > high))
        mask = mlow + mhigh
        return mask
//  [python]
def _compute_window_mask(self, max_alignments):
        """Compute window mask for inference.
        Args:
            max_alignments (int): [batch_size]
        """
        expanded_max_alignments = tf.expand_dims(max_alignments, 1)  # [batch_size, 1]
        low = expanded_max_alignments - self.win_front
        high = expanded_max_alignments + self.win_back
        mlow = tf.cast((self._indices < low), tf.float32)
        mhigh = tf.cast((self._indices > high), tf.float32)
        mask = mlow + mhigh
        return mask
//  [python]
def losses(logits, labels):
    with tf.variable_scope('loss') as scope:
        
        labels = tf.to_int64(labels)
        
        # to use this loss fuction, one-hot encoding is needed!
        cross_entropy = tf.nn.softmax_cross_entropy_with_logits\
                        (logits=logits, labels=labels, name='xentropy_per_example')

                        
#        cross_entropy = tf.nn.sparse_softmax_cross_entropy_with_logits\
#                        (logits=logits, labels=labels, name='xentropy_per_example')
                        
        loss = tf.reduce_mean(cross_entropy, name='loss')
        tf.summary.scalar(scope.name+'/loss', loss)
        
    return loss
//  [python]
def losses(logits, labels):
    with tf.variable_scope('loss') as scope:
        
        labels = tf.cast(labels, tf.int64)
        
        # to use this loss fuction, one-hot encoding is needed!
        cross_entropy = tf.nn.softmax_cross_entropy_with_logits\
                        (logits=logits, labels=labels, name='xentropy_per_example')

                        
#        cross_entropy = tf.nn.sparse_softmax_cross_entropy_with_logits\
#                        (logits=logits, labels=labels, name='xentropy_per_example')
                        
        loss = tf.reduce_mean(cross_entropy, name='loss')
        tf.summary.scalar(scope.name+'/loss', loss)
        
    return loss
//  [python]
def evaluate_accuracy(data_iter, net):
    acc_sum, n = 0.0, 0
    for _, (X, y) in enumerate(data_iter):
        y = t.cast(y,dtype=t.int64)
        acc_sum += np.sum(t.cast(t.argmax(net(X), axis=1), dtype=t.int64) == y)
        n += y.shape[0]
    return acc_sum / n
//  [python]
def evaluate_accuracy(data_iter, net):
    acc_sum, n = 0.0, 0
    for _, (X, y) in enumerate(data_iter):
        y = tf.cast(y,dtype=tf.int64)
        acc_sum += np.sum(tf.cast(tf.argmax(net(X), axis=1), dtype=tf.int64) == y)
        n += y.shape[0]
    return acc_sum / n
//  [python]
def extract_single_agent_trajectory(self, agent_id, trajectory):
    """Pull a single agent's experience out of a trajectory and makes a new one.

    Note: trajectory format (B = batch size, T = # timesteps, N = # agents)
      observation: DictWrapper. Need to pull out agent's observations.
      step_type: (B, T). Agent's is the same as global step type.
      next_step_type: (B, T). Agent's is the same as global step type.
      action: (B, T, N). Need to pull out individual agent's action
      policy_info: ListWrapper of length N containing logits for each agent.
      reward: (B, T, N). Need to pull out individual agent's rewards.
      discount: (B, T).

    Args:
      agent_id: Integer ID for this agent.
      trajectory: As above

    Returns:
      A new trajectory for this agent only
    """
    agent_trajectory = super().extract_single_agent_trajectory(agent_id,
                                                               trajectory)
    reward = trajectory.reward[:, :, agent_id]
    policy_info = trajectory.policy_info[agent_id]
    if 'attention_weights' in policy_info:
      p = policy_info['attention_weights']
      eps = tf.keras.backend.epsilon()
      p += eps
      raw_bonus = 0
      score_multiplier = t.cast(
          tf.clip_by_value(self._global_step, 0, self._bonus_timescale) /
          self._bonus_timescale * self._bonus_ratio, t.float32)
      for j, other_policy_info in enumerate(trajectory.policy_info):
        if j == agent_id:
          continue
        q = other_policy_info['attention_weights']
        q += eps
        if self.attention_bonus_type.lower() == 'kld':
          raw_bonus -= tf.reduce_sum(p * tf.math.log(p / q), axis=(-1, -2))
        elif self.attention_bonus_type.lower() == 'jsd':
          m = (p + q) / 2
          raw_bonus -= (
              tf.reduce_sum(p * tf.math.log(p / m), axis=(-1, -2)) +
              tf.reduce_sum(q * tf.math.log(q / m), axis=(-1, -2))) / 2
        else:
          raise NotImplementedError
      scaled_bonus = raw_bonus * score_multiplier
      reward += scaled_bonus
      if self._unscaled_bonus_metric:
        self._unscaled_bonus_metric(raw_bonus, agent_id)
      if self._scaled_bonus_metric:
        self._scaled_bonus_metric(scaled_bonus, agent_id)

    agent_trajectory.replace(reward=reward)

    return agent_trajectory
//  [python]
def extract_single_agent_trajectory(self, agent_id, trajectory):
    """Pull a single agent's experience out of a trajectory and makes a new one.

    Note: trajectory format (B = batch size, T = # timesteps, N = # agents)
      observation: DictWrapper. Need to pull out agent's observations.
      step_type: (B, T). Agent's is the same as global step type.
      next_step_type: (B, T). Agent's is the same as global step type.
      action: (B, T, N). Need to pull out individual agent's action
      policy_info: ListWrapper of length N containing logits for each agent.
      reward: (B, T, N). Need to pull out individual agent's rewards.
      discount: (B, T).

    Args:
      agent_id: Integer ID for this agent.
      trajectory: As above

    Returns:
      A new trajectory for this agent only
    """
    agent_trajectory = super().extract_single_agent_trajectory(agent_id,
                                                               trajectory)
    reward = trajectory.reward[:, :, agent_id]
    policy_info = trajectory.policy_info[agent_id]
    if 'attention_weights' in policy_info:
      p = policy_info['attention_weights']
      eps = tf.keras.backend.epsilon()
      p += eps
      raw_bonus = 0
      score_multiplier = tf.cast(
          tf.clip_by_value(self._global_step, 0, self._bonus_timescale) /
          self._bonus_timescale * self._bonus_ratio, tf.float32)
      for j, other_policy_info in enumerate(trajectory.policy_info):
        if j == agent_id:
          continue
        q = other_policy_info['attention_weights']
        q += eps
        if self.attention_bonus_type.lower() == 'kld':
          raw_bonus -= tf.reduce_sum(p * tf.math.log(p / q), axis=(-1, -2))
        elif self.attention_bonus_type.lower() == 'jsd':
          m = (p + q) / 2
          raw_bonus -= (
              tf.reduce_sum(p * tf.math.log(p / m), axis=(-1, -2)) +
              tf.reduce_sum(q * tf.math.log(q / m), axis=(-1, -2))) / 2
        else:
          raise NotImplementedError
      scaled_bonus = raw_bonus * score_multiplier
      reward += scaled_bonus
      if self._unscaled_bonus_metric:
        self._unscaled_bonus_metric(raw_bonus, agent_id)
      if self._scaled_bonus_metric:
        self._scaled_bonus_metric(scaled_bonus, agent_id)

    agent_trajectory.replace(reward=reward)

    return agent_trajectory
//  [python]
def voc_label_indices(colormap, colormap2label):
    """
    convert colormap (tf image) to colormap2label (uint8 tensor).
    """
    colormap = t.cast(colormap, dtype=tf.int32)
    idx = tf.add(tf.multiply(colormap[:, :, 0], 256), colormap[:, :, 1])
    idx = tf.add(tf.multiply(idx, 256), colormap[:, :, 2])

    return tf.gather_nd(colormap2label, tf.expand_dims(idx, -1))
//  [python]
def voc_label_indices(colormap, colormap2label):
    """
    convert colormap (tf image) to colormap2label (uint8 tensor).
    """
    colormap = tf.cast(colormap, dtype=tf.int32)
    idx = tf.add(tf.multiply(colormap[:, :, 0], 256), colormap[:, :, 1])
    idx = tf.add(tf.multiply(idx, 256), colormap[:, :, 2])

    return tf.gather_nd(colormap2label, tf.expand_dims(idx, -1))
//  [python]
def call(self, inputs, training=False):
        hidden_states, attention_mask, head_mask = inputs

        batch_size = shape_list(hidden_states)[0]
        mixed_query_layer = self.query(hidden_states)
        mixed_key_layer = self.key(hidden_states)
        mixed_value_layer = self.value(hidden_states)

        query_layer = self.transpose_for_scores(mixed_query_layer, batch_size)
        key_layer = self.transpose_for_scores(mixed_key_layer, batch_size)
        value_layer = self.transpose_for_scores(mixed_value_layer, batch_size)

        # Take the dot product between "query" and "key" to get the raw attention scores.
        # (batch size, num_heads, seq_len_q, seq_len_k)
        attention_scores = tf.matmul(query_layer, key_layer, transpose_b=True)
        # scale attention_scores
        dk = tf.to_float(shape_list(key_layer)[-1])
        attention_scores = attention_scores / tf.math.sqrt(dk)

        if attention_mask is not None:
            # Apply the attention mask is (precomputed for all layers in TFAlbertModel call() function)
            attention_scores = attention_scores + attention_mask

        # Normalize the attention scores to probabilities.
        attention_probs = tf.nn.softmax(attention_scores, axis=-1)

        # This is actually dropping out entire tokens to attend to, which might
        # seem a bit unusual, but is taken from the original Transformer paper.
        attention_probs = self.dropout(attention_probs, training=training)

        # Mask heads if we want to
        if head_mask is not None:
            attention_probs = attention_probs * head_mask

        context_layer = tf.matmul(attention_probs, value_layer)

        context_layer = tf.transpose(context_layer, perm=[0, 2, 1, 3])
        context_layer = tf.reshape(
            context_layer, (batch_size, -1, self.all_head_size)
        )  # (batch_size, seq_len_q, all_head_size)

        outputs = (context_layer, attention_probs) if self.output_attentions else (context_layer,)
        return outputs
//  [python]
def call(self, inputs, training=False):
        hidden_states, attention_mask, head_mask = inputs

        batch_size = shape_list(hidden_states)[0]
        mixed_query_layer = self.query(hidden_states)
        mixed_key_layer = self.key(hidden_states)
        mixed_value_layer = self.value(hidden_states)

        query_layer = self.transpose_for_scores(mixed_query_layer, batch_size)
        key_layer = self.transpose_for_scores(mixed_key_layer, batch_size)
        value_layer = self.transpose_for_scores(mixed_value_layer, batch_size)

        # Take the dot product between "query" and "key" to get the raw attention scores.
        # (batch size, num_heads, seq_len_q, seq_len_k)
        attention_scores = tf.matmul(query_layer, key_layer, transpose_b=True)
        # scale attention_scores
        dk = tf.cast(shape_list(key_layer)[-1], tf.float32)
        attention_scores = attention_scores / tf.math.sqrt(dk)

        if attention_mask is not None:
            # Apply the attention mask is (precomputed for all layers in TFAlbertModel call() function)
            attention_scores = attention_scores + attention_mask

        # Normalize the attention scores to probabilities.
        attention_probs = tf.nn.softmax(attention_scores, axis=-1)

        # This is actually dropping out entire tokens to attend to, which might
        # seem a bit unusual, but is taken from the original Transformer paper.
        attention_probs = self.dropout(attention_probs, training=training)

        # Mask heads if we want to
        if head_mask is not None:
            attention_probs = attention_probs * head_mask

        context_layer = tf.matmul(attention_probs, value_layer)

        context_layer = tf.transpose(context_layer, perm=[0, 2, 1, 3])
        context_layer = tf.reshape(
            context_layer, (batch_size, -1, self.all_head_size)
        )  # (batch_size, seq_len_q, all_head_size)

        outputs = (context_layer, attention_probs) if self.output_attentions else (context_layer,)
        return outputs
//  [python]
def bias_func(x, bias_prob=0.5):
    # Bias-conditioning function of choice.
    x = tf.exp(x)

    # Randomized application.
    if bias_prob < 1:
        bias_bit = tf.less(tf.random.uniform(shape=[]), bias_prob)
        bias_not = tf.logical_not(bias_bit)
        x = x * t.cast(bias_bit, x.dtype) + t.cast(bias_not, x.dtype)

    return x
//  [python]
def bias_func(x, bias_prob=0.5):
    # Bias-conditioning function of choice.
    x = tf.exp(x)

    # Randomized application.
    if bias_prob < 1:
        bias_bit = tf.less(tf.random.uniform(shape=[]), bias_prob)
        bias_not = tf.logical_not(bias_bit)
        x = x * tf.cast(bias_bit, x.dtype) + tf.cast(bias_not, x.dtype)

    return x
//  [python]
def train_ch3(net, train_iter, test_iter, loss, num_epochs, batch_size, params=None, lr=None, trainer=None):
    for epoch in range(num_epochs):
        train_l_sum, train_acc_sum, n = 0.0, 0.0, 0
        for X, y in train_iter:
            with tf.GradientTape() as tape:
                y_hat = net(X)
                l = tf.reduce_sum(loss(y_hat, y))
            grads = tape.gradient(l, params)
            if trainer is None:
                # 如果没有传入优化器，则使用原先编写的小批量随机梯度下降
                for i, param in enumerate(params):
                    param.assign_sub(lr * grads[i] / batch_size)
            else:
                # tf.keras.optimizers.SGD 直接使用是随机梯度下降 theta(t+1) = theta(t) - learning_rate * gradient
                # 这里使用批量梯度下降，需要对梯度除以 batch_size, 对应原书代码的 trainer.step(batch_size)
                trainer.apply_gradients(zip([grad / batch_size for grad in grads], params))  
                
            y = tf.cast(y, dtype=tf.float32)
            train_l_sum += l.numpy()
            train_acc_sum += tf.reduce_sum(tf.cast(tf.argmax(y_hat, axis=1) == tf.cast(y, dtype=tf.int64), dtype=tf.int64)).numpy()
            n += y.shape[0]
        test_acc = evaluate_accuracy(test_iter, net)
        print('epoch %d, loss %.4f, train acc %.3f, test acc %.3f'% (epoch + 1, train_l_sum / n, train_acc_sum / n, test_acc))
//  [python]
def call(self, inputs, training=False):
        input_tensor, attention_mask, head_mask = inputs

        batch_size = shape_list(input_tensor)[0]
        mixed_query_layer = self.query(input_tensor)
        mixed_key_layer = self.key(input_tensor)
        mixed_value_layer = self.value(input_tensor)

        query_layer = self.transpose_for_scores(mixed_query_layer, batch_size)
        key_layer = self.transpose_for_scores(mixed_key_layer, batch_size)
        value_layer = self.transpose_for_scores(mixed_value_layer, batch_size)

        # Take the dot product between "query" and "key" to get the raw attention scores.
        # (batch size, num_heads, seq_len_q, seq_len_k)
        attention_scores = tf.matmul(query_layer, key_layer, transpose_b=True)
        # scale attention_scores
        dk = tf.to_float(shape_list(key_layer)[-1])
        attention_scores = attention_scores / tf.math.sqrt(dk)

        if attention_mask is not None:
            # Apply the attention mask is (precomputed for all layers in TFBertModel call() function)
            attention_scores = attention_scores + attention_mask

        # Normalize the attention scores to probabilities.
        attention_probs = tf.nn.softmax(attention_scores, axis=-1)

        # This is actually dropping out entire tokens to attend to, which might
        # seem a bit unusual, but is taken from the original Transformer paper.
        attention_probs = self.dropout(attention_probs, training=training)

        # Mask heads if we want to
        if head_mask is not None:
            attention_probs = attention_probs * head_mask

        context_layer = tf.matmul(attention_probs, value_layer)

        context_layer = tf.transpose(context_layer, perm=[0, 2, 1, 3])
        context_layer = tf.reshape(
            context_layer, (batch_size, -1, self.all_head_size)
        )  # (batch_size, seq_len_q, all_head_size)

        self_outputs = (context_layer, attention_probs) if self.output_attentions else (context_layer,)

        hidden_states = self_outputs[0]

        hidden_states = self.dense(hidden_states)
        hidden_states = self.dropout(hidden_states, training=training)
        attention_output = self.LayerNorm(hidden_states + input_tensor)

        # add attentions if we output them
        outputs = (attention_output,) + self_outputs[1:]
        return outputs
//  [python]
def call(self, inputs, training=False):
        input_tensor, attention_mask, head_mask = inputs

        batch_size = shape_list(input_tensor)[0]
        mixed_query_layer = self.query(input_tensor)
        mixed_key_layer = self.key(input_tensor)
        mixed_value_layer = self.value(input_tensor)

        query_layer = self.transpose_for_scores(mixed_query_layer, batch_size)
        key_layer = self.transpose_for_scores(mixed_key_layer, batch_size)
        value_layer = self.transpose_for_scores(mixed_value_layer, batch_size)

        # Take the dot product between "query" and "key" to get the raw attention scores.
        # (batch size, num_heads, seq_len_q, seq_len_k)
        attention_scores = tf.matmul(query_layer, key_layer, transpose_b=True)
        # scale attention_scores
        dk = tf.cast(shape_list(key_layer)[-1], tf.float32)
        attention_scores = attention_scores / tf.math.sqrt(dk)

        if attention_mask is not None:
            # Apply the attention mask is (precomputed for all layers in TFBertModel call() function)
            attention_scores = attention_scores + attention_mask

        # Normalize the attention scores to probabilities.
        attention_probs = tf.nn.softmax(attention_scores, axis=-1)

        # This is actually dropping out entire tokens to attend to, which might
        # seem a bit unusual, but is taken from the original Transformer paper.
        attention_probs = self.dropout(attention_probs, training=training)

        # Mask heads if we want to
        if head_mask is not None:
            attention_probs = attention_probs * head_mask

        context_layer = tf.matmul(attention_probs, value_layer)

        context_layer = tf.transpose(context_layer, perm=[0, 2, 1, 3])
        context_layer = tf.reshape(
            context_layer, (batch_size, -1, self.all_head_size)
        )  # (batch_size, seq_len_q, all_head_size)

        self_outputs = (context_layer, attention_probs) if self.output_attentions else (context_layer,)

        hidden_states = self_outputs[0]

        hidden_states = self.dense(hidden_states)
        hidden_states = self.dropout(hidden_states, training=training)
        attention_output = self.LayerNorm(hidden_states + input_tensor)

        # add attentions if we output them
        outputs = (attention_output,) + self_outputs[1:]
        return outputs