jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def equalize(image):
  """Implements Equalize function from PIL using TF ops."""

  def scale_channel(im, c):
    """Scale the data in the channel to implement equalize."""
    im = tf.cast(im[:, :, c], tf.int32)
    # Compute the histogram of the image channel.
    histo = tf.histogram_fixed_width(im, [0, 255], nbins=256)

    # For the purposes of computing the step, filter out the nonzeros.
    nonzero = tf.where(tf.not_equal(histo, 0))
    nonzero_histo = tf.reshape(tf.gather(histo, nonzero), [-1])
    step = (tf.reduce_sum(nonzero_histo) - nonzero_histo[-1]) // 255

    def build_lut(histo, step):
      # Compute the cumulative sum, shifting by step // 2
      # and then normalization by step.
      lut = (tf.cumsum(histo) + (step // 2)) // step
      # Shift lut, prepending with 0.
      lut = tf.concat([[0], lut[:-1]], 0)
      # Clip the counts to be in range.  This is done
      # in the C code for image.point.
      return tf.clip_by_value(lut, 0, 255)

    # If step is zero, return the original image.  Otherwise, build
    # lut from the full histogram and step and then index from it.
    result = tf.cond(
        tf.equal(step, 0), lambda: im,
        lambda: tf.gather(build_lut(histo, step), im))

    return tf.cast(result, tf.uint8)

  # Assumes RGB for now.  Scales each channel independently
  # and then stacks the result.
  s1 = scale_channel(image, 0)
  s2 = scale_channel(image, 1)
  s3 = scale_channel(image, 2)
  image = tf.stack([s1, s2, s3], 2)
  return image
//  [python]
def random_erase(image,
                 prob,
                 min_area=0.02,
                 max_area=1 / 3,
                 min_aspect=1 / 3,
                 max_aspect=10 / 3,
                 mode='pixel'):
  """The random erasing augmentations: https://arxiv.org/pdf/1708.04896.pdf.

  This augmentation is applied after image normalization.

  Args:
    image: Input image after all other augmentation and normalization. It has
      type tf.float32.
    prob: Probability of applying the random erasing operation.
    min_area: As named.
    max_area: As named.
    min_aspect: As named.
    max_aspect: As named.
    mode: How the erased area is filled. 'pixel' means white noise (uniform
      dist).

  Returns:
    Randomly erased image.
  """

  image_height = tf.shape(image)[0]
  image_width = tf.shape(image)[1]
  image_area = tf.cast(image_width * image_height, tf.float32)

  # Sample width, height
  erase_area = tf.random.uniform([], min_area, max_area) * image_area
  log_max_target_ar = tf.math.log(
      tf.minimum(
          tf.math.divide(
              tf.math.square(tf.cast(image_width, tf.float32)), erase_area),
          max_aspect))
  log_min_target_ar = tf.math.log(
      tf.maximum(
          tf.math.divide(erase_area,
                         tf.math.square(tf.cast(image_height, tf.float32))),
          min_aspect))
  erase_aspect_ratio = tf.math.exp(
      tf.random.uniform([], log_min_target_ar, log_max_target_ar))
  erase_h = tf.cast(tf.math.sqrt(erase_area / erase_aspect_ratio), tf.int32)
  erase_w = tf.cast(tf.math.sqrt(erase_area * erase_aspect_ratio), tf.int32)

  # Sample (left, top) of the rectangle to erase
  erase_left = tf.random.uniform(
      shape=[], minval=0, maxval=image_width - erase_w, dtype=tf.int32)
  erase_top = tf.random.uniform(
      shape=[], minval=0, maxval=image_height - erase_h, dtype=tf.int32)
  pad_right = image_width - erase_w - erase_left
  pad_bottom = image_height - erase_h - erase_top
  mask = tf.pad(
      tf.zeros([erase_h, erase_w], dtype=image.dtype),
      [[erase_top, pad_bottom], [erase_left, pad_right]],
      constant_values=1)
  mask = tf.expand_dims(mask, -1)  # [H, W, 1]
  if mode == 'pixel':
    fill = tf.random.truncated_normal(
        tf.shape(image), 0.0, 1.0, dtype=image.dtype)
  else:
    fill = tf.zeros(tf.shape(image), dtype=image.dtype)

  should_apply_op = tf.floor(tf.random.uniform([], dtype=tf.float32) + prob) > 0
  augmented_image = tf.cond(should_apply_op,
                            lambda: mask * image + (1 - mask) * fill,
                            lambda: image)
  return augmented_image
//  [python]
def random_erase(image,
                 prob,
                 min_area=0.02,
                 max_area=1 / 3,
                 min_aspect=1 / 3,
                 max_aspect=10 / 3,
                 mode='pixel'):
  """The random erasing augmentations: https://arxiv.org/pdf/1708.04896.pdf.

  This augmentation is applied after image normalization.

  Args:
    image: Input image after all other augmentation and normalization. It has
      type tf.float32.
    prob: Probability of applying the random erasing operation.
    min_area: As named.
    max_area: As named.
    min_aspect: As named.
    max_aspect: As named.
    mode: How the erased area is filled. 'pixel' means white noise (uniform
      dist).

  Returns:
    Randomly erased image.
  """

  image_height = tf.shape(image)[0]
  image_width = tf.shape(image)[1]
  image_area = tf.cast(image_width * image_height, tf.float32)

  # Sample width, height
  erase_area = tf.random.uniform([], min_area, max_area) * image_area
  log_max_target_ar = tf.math.log(
      tf.minimum(
          tf.math.divide(
              tf.math.square(tf.cast(image_width, tf.float32)), erase_area),
          max_aspect))
  log_min_target_ar = tf.math.log(
      tf.maximum(
          tf.math.divide(erase_area,
                         tf.math.square(tf.cast(image_height, tf.float32))),
          min_aspect))
  erase_aspect_ratio = tf.math.exp(
      tf.random.uniform([], log_min_target_ar, log_max_target_ar))
  erase_h = tf.cast(tf.math.sqrt(erase_area / erase_aspect_ratio), tf.int32)
  erase_w = tf.cast(tf.math.sqrt(erase_area * erase_aspect_ratio), tf.int32)

  # Sample (left, top) of the rectangle to erase
  erase_left = tf.random.uniform(
      shape=[], minval=0, maxval=image_width - erase_w, dtype=tf.int32)
  erase_top = tf.random.uniform(
      shape=[], minval=0, maxval=image_height - erase_h, dtype=tf.int32)
  pad_right = image_width - erase_w - erase_left
  pad_bottom = image_height - erase_h - erase_top
  mask = tf.pad(
      tf.zeros([erase_h, erase_w], dtype=image.dtype),
      [[erase_top, pad_bottom], [erase_left, pad_right]],
      constant_values=1)
  mask = tf.expand_dims(mask, -1)  # [H, W, 1]
  if mode == 'pixel':
    fill = tf.random.truncated_normal(
        tf.shape(image), 0.0, 1.0, dtype=image.dtype)
  else:
    fill = tf.zeros(tf.shape(image), dtype=image.dtype)

  should_apply_op = tf.cast(
      tf.floor(tf.random.uniform([], dtype=tf.float32) + prob), tf.bool)
  augmented_image = tf.cond(should_apply_op,
                            lambda: mask * image + (1 - mask) * fill,
                            lambda: image)
  return augmented_image
//  [python]
def _apply_func_with_prob(func, image, args, prob):
  """Apply `func` to image w/ `args` as input with probability `prob`."""
  assert isinstance(args, tuple)

  # Apply the function with probability `prob`.
  should_apply_op = tf.floor(tf.random.uniform([], dtype=tf.float32) + prob) > 0
  augmented_image = tf.cond(should_apply_op, lambda: func(image, *args),
                            lambda: image)
  return augmented_image
//  [python]
def _apply_func_with_prob(func, image, args, prob):
  """Apply `func` to image w/ `args` as input with probability `prob`."""
  assert isinstance(args, tuple)

  # Apply the function with probability `prob`.
  should_apply_op = tf.cast(
      tf.floor(tf.random.uniform([], dtype=tf.float32) + prob), tf.bool)
  augmented_image = tf.cond(should_apply_op, lambda: func(image, *args),
                            lambda: image)
  return augmented_image
//  [python]
def train_decode(self, decode_outputs):
    logits = embedding_linear(self.decoder_layer.embedding_lookup.embeddings,
                              decode_outputs)
    decode_output_ids = tf.to_int32(tf.argmax(logits, axis=-1))
    output_log_probs = tf.nn.log_softmax(logits, axis=-1)
    return logits, decode_output_ids, output_log_probs
//  [python]
def train_decode(self, decode_outputs):
    logits = embedding_linear(self.decoder_layer.embedding_lookup.embeddings,
                              decode_outputs)
    decode_output_ids = tf.cast(tf.argmax(logits, axis=-1), tf.int32)
    output_log_probs = tf.nn.log_softmax(logits, axis=-1)
    return logits, decode_output_ids, output_log_probs
//  [python]
def __call__(
      self,
      y_true: tf.Tensor,
      y_pred: tf.Tensor,
      sample_weight: Optional[tf.Tensor] = None,
  ) -> tf.Tensor:
    y_true = tf.to_int32(tf.reshape(y_true, [-1]))
    y_true_one_hot = tf.one_hot(y_true, self._num_classes)
    return super().__call__(y_true_one_hot, y_pred, sample_weight=sample_weight)
//  [python]
def __call__(
      self,
      y_true: tf.Tensor,
      y_pred: tf.Tensor,
      sample_weight: Optional[tf.Tensor] = None,
  ) -> tf.Tensor:
    y_true = tf.cast(tf.reshape(y_true, [-1]), tf.int32)
    y_true_one_hot = tf.one_hot(y_true, self._num_classes)
    return super().__call__(y_true_one_hot, y_pred, sample_weight=sample_weight)
//  [python]
def bilinear_sampler(imgs, coords):
	"""
	Construct a new image by bilinear sampling from the input image.
	Points falling outside the source image boundary have value 0.
	Args:
		imgs: source image to be sampled from [batch, height_s, width_s, channels]
		coords: coordinates of source pixels to sample from [batch, height_t,width_t, 2]. height_t/width_t correspond to the dimensions of the outputimage (don't need to be the same as height_s/width_s). The two channels correspond to x and y coordinates respectively.
	Returns:
		A new sampled image [batch, height_t, width_t, channels]
	"""

	def _repeat(x, n_repeats):
		rep = tf.transpose(
			tf.expand_dims(tf.ones(shape=tf.stack([
				n_repeats,
			])), 1), [1, 0])
		rep = tf.cast(rep, 'float32')
		x = tf.matmul(tf.reshape(x, (-1, 1)), rep)
		return tf.reshape(x, [-1])

	with tf.name_scope('image_sampling'):
		coords_x, coords_y = tf.split(coords, [1, 1], axis=3)
		inp_size = tf.shape(imgs)
		coord_size = tf.shape(coords)
		out_size = [coord_size[0],coord_size[1],coord_size[2],inp_size[3]]

		coords_x = tf.to_float(coords_x)
		coords_y = tf.to_float(coords_y)

		x0 = tf.floor(coords_x)
		x1 = x0 + 1
		y0 = tf.floor(coords_y)
		y1 = y0 + 1

		y_max = tf.to_float(inp_size[1] - 1)
		x_max = tf.to_float(inp_size[2] - 1)
		zero = tf.zeros([1], dtype='float32')

		wt_x0 = x1 - coords_x
		wt_x1 = coords_x - x0
		wt_y0 = y1 - coords_y
		wt_y1 = coords_y - y0

		x0_safe = tf.clip_by_value(x0, zero[0], x_max)
		y0_safe = tf.clip_by_value(y0, zero[0], y_max)
		x1_safe = tf.clip_by_value(x1, zero[0], x_max)
		y1_safe = tf.clip_by_value(y1, zero[0], y_max)

		## indices in the flat image to sample from
		dim2 = tf.to_float(inp_size[2])
		dim1 = tf.to_float(inp_size[2] * inp_size[1])
		base = tf.reshape(_repeat(tf.to_float(tf.range(coord_size[0])) * dim1,coord_size[1] * coord_size[2]),[out_size[0], out_size[1], out_size[2], 1])

		base_y0 = base + y0_safe * dim2
		base_y1 = base + y1_safe * dim2
		idx00 = x0_safe + base_y0
		idx01 = x0_safe + base_y1
		idx10 = x1_safe + base_y0
		idx11 = x1_safe + base_y1

		## sample from imgs
		imgs_flat = tf.reshape(imgs, tf.stack([-1, inp_size[3]]))
		imgs_flat = tf.to_float(imgs_flat)
		im00 = tf.reshape(tf.gather(imgs_flat, tf.cast(idx00, 'int32')), out_size)
		im01 = tf.reshape(tf.gather(imgs_flat, tf.cast(idx01, 'int32')), out_size)
		im10 = tf.reshape(tf.gather(imgs_flat, tf.cast(idx10, 'int32')), out_size)
		im11 = tf.reshape(tf.gather(imgs_flat, tf.cast(idx11, 'int32')), out_size)

		w00 = wt_x0 * wt_y0
		w01 = wt_x0 * wt_y1
		w10 = wt_x1 * wt_y0
		w11 = wt_x1 * wt_y1

		output = tf.add_n([
			w00 * im00, w01 * im01,
			w10 * im10, w11 * im11
		])

		return output
//  [python]
def bilinear_sampler(imgs, coords):
	"""
	Construct a new image by bilinear sampling from the input image.
	Points falling outside the source image boundary have value 0.
	Args:
		imgs: source image to be sampled from [batch, height_s, width_s, channels]
		coords: coordinates of source pixels to sample from [batch, height_t,width_t, 2]. height_t/width_t correspond to the dimensions of the outputimage (don't need to be the same as height_s/width_s). The two channels correspond to x and y coordinates respectively.
	Returns:
		A new sampled image [batch, height_t, width_t, channels]
	"""

	def _repeat(x, n_repeats):
		rep = tf.transpose(
			tf.expand_dims(tf.ones(shape=tf.stack([
				n_repeats,
			])), 1), [1, 0])
		rep = tf.cast(rep, 'float32')
		x = tf.matmul(tf.reshape(x, (-1, 1)), rep)
		return tf.reshape(x, [-1])

	with tf.name_scope('image_sampling'):
		coords_x, coords_y = tf.split(coords, [1, 1], axis=3)
		inp_size = tf.shape(imgs)
		coord_size = tf.shape(coords)
		out_size = [coord_size[0],coord_size[1],coord_size[2],inp_size[3]]

		coords_x = tf.cast(coords_x, 'float32')
		coords_y = tf.cast(coords_y, 'float32')

		x0 = tf.floor(coords_x)
		x1 = x0 + 1
		y0 = tf.floor(coords_y)
		y1 = y0 + 1

		y_max = tf.cast(inp_size[1] - 1, 'float32')
		x_max = tf.cast(inp_size[2] - 1, 'float32')
		zero = tf.zeros([1], dtype='float32')

		wt_x0 = x1 - coords_x
		wt_x1 = coords_x - x0
		wt_y0 = y1 - coords_y
		wt_y1 = coords_y - y0

		x0_safe = tf.clip_by_value(x0, zero[0], x_max)
		y0_safe = tf.clip_by_value(y0, zero[0], y_max)
		x1_safe = tf.clip_by_value(x1, zero[0], x_max)
		y1_safe = tf.clip_by_value(y1, zero[0], y_max)

		## indices in the flat image to sample from
		dim2 = tf.cast(inp_size[2], 'float32')
		dim1 = tf.cast(inp_size[2] * inp_size[1], 'float32')
		base = tf.reshape(_repeat(tf.cast(tf.range(coord_size[0]), 'float32') * dim1,coord_size[1] * coord_size[2]),[out_size[0], out_size[1], out_size[2], 1])

		base_y0 = base + y0_safe * dim2
		base_y1 = base + y1_safe * dim2
		idx00 = x0_safe + base_y0
		idx01 = x0_safe + base_y1
		idx10 = x1_safe + base_y0
		idx11 = x1_safe + base_y1

		## sample from imgs
		imgs_flat = tf.reshape(imgs, tf.stack([-1, inp_size[3]]))
		imgs_flat = tf.cast(imgs_flat, 'float32')
		im00 = tf.reshape(tf.gather(imgs_flat, tf.cast(idx00, 'int32')), out_size)
		im01 = tf.reshape(tf.gather(imgs_flat, tf.cast(idx01, 'int32')), out_size)
		im10 = tf.reshape(tf.gather(imgs_flat, tf.cast(idx10, 'int32')), out_size)
		im11 = tf.reshape(tf.gather(imgs_flat, tf.cast(idx11, 'int32')), out_size)

		w00 = wt_x0 * wt_y0
		w01 = wt_x0 * wt_y1
		w10 = wt_x1 * wt_y0
		w11 = wt_x1 * wt_y1

		output = tf.add_n([
			w00 * im00, w01 * im01,
			w10 * im10, w11 * im11
		])

		return output
//  [python]
def normalize(image, label):
    """Convert `image` from [0, 255] -> [-0.5, 0.5] floats."""
    image = tf.cast(image, tf.float32) * (1.0 / 255) - 0.5
    return image, label
//  [python]
def __call__(self,
               y_true: tf.Tensor,
               y_pred: tf.Tensor,
               sample_weight: Optional[tf.Tensor] = None) -> tf.Tensor:
    if self._class_weight:
      class_weight = tf.convert_to_tensor(self._class_weight, dtype=tf.float32)
      label = tf.argmax(y_true, axis=1)
      loss_weight = tf.gather(class_weight, label)
    else:
      loss_weight = tf.ones(tf.shape(y_true)[0])
    y_true = tf.to_float(y_true)
    y_pred = tf.clip_by_value(y_pred, self._epsilon, 1 - self._epsilon)
    batch_size = tf.cast(tf.shape(y_pred)[0], y_pred.dtype)
    if sample_weight is None:
      sample_weight = tf.constant(1.0)
    weight_shape = sample_weight.shape
    weight_rank = weight_shape.ndims
    y_pred_rank = y_pred.shape.ndims
    if y_pred_rank - weight_rank == 1:
      sample_weight = tf.expand_dims(sample_weight, [-1])
    elif weight_rank != 0:
      raise ValueError(f'Unexpected sample_weights, should be either a scalar'
                       f'or a vector of batch_size:{batch_size.numpy()}')
    ce = -tf.math.log(y_pred)
    modulating_factor = tf.math.pow(1 - y_pred, self._gamma)
    losses = y_true * modulating_factor * ce * sample_weight
    losses = losses * loss_weight[:, tf.newaxis]
    # By default, this function uses "sum_over_batch_size" reduction for the
    # loss per batch.
    return tf.reduce_sum(losses) / batch_size
//  [python]
def __call__(self,
               y_true: tf.Tensor,
               y_pred: tf.Tensor,
               sample_weight: Optional[tf.Tensor] = None) -> tf.Tensor:
    if self._class_weight:
      class_weight = tf.convert_to_tensor(self._class_weight, dtype=tf.float32)
      label = tf.argmax(y_true, axis=1)
      loss_weight = tf.gather(class_weight, label)
    else:
      loss_weight = tf.ones(tf.shape(y_true)[0])
    y_true = tf.cast(y_true, y_pred.dtype)
    y_pred = tf.clip_by_value(y_pred, self._epsilon, 1 - self._epsilon)
    batch_size = tf.cast(tf.shape(y_pred)[0], y_pred.dtype)
    if sample_weight is None:
      sample_weight = tf.constant(1.0)
    weight_shape = sample_weight.shape
    weight_rank = weight_shape.ndims
    y_pred_rank = y_pred.shape.ndims
    if y_pred_rank - weight_rank == 1:
      sample_weight = tf.expand_dims(sample_weight, [-1])
    elif weight_rank != 0:
      raise ValueError(f'Unexpected sample_weights, should be either a scalar'
                       f'or a vector of batch_size:{batch_size.numpy()}')
    ce = -tf.math.log(y_pred)
    modulating_factor = tf.math.pow(1 - y_pred, self._gamma)
    losses = y_true * modulating_factor * ce * sample_weight
    losses = losses * loss_weight[:, tf.newaxis]
    # By default, this function uses "sum_over_batch_size" reduction for the
    # loss per batch.
    return tf.reduce_sum(losses) / batch_size
//  [python]
def decode(serialized_example):
    """Parses an image and label from the given `serialized_example`."""
    features = tf.parse_single_example(
        serialized_example,
        # Defaults are not specified since both keys are required.
        features={
            "image_raw": tf.FixedLenFeature([], tf.string),
            "label": tf.FixedLenFeature([], tf.int64),
        },
    )

    # Convert from a scalar string tensor (whose single string has
    # length mnist.IMAGE_PIXELS) to a uint8 tensor with shape
    # [mnist.IMAGE_PIXELS].
    image = tf.decode_raw(features["image_raw"], tf.uint8)
    image.set_shape((mnist.IMAGE_PIXELS))

    # Convert label from a scalar uint8 tensor to an int32 scalar.
    label = tf.to_int32(features["label"])

    return image, label
//  [python]
def decode(serialized_example):
    """Parses an image and label from the given `serialized_example`."""
    features = tf.parse_single_example(
        serialized_example,
        # Defaults are not specified since both keys are required.
        features={
            "image_raw": tf.FixedLenFeature([], tf.string),
            "label": tf.FixedLenFeature([], tf.int64),
        },
    )

    # Convert from a scalar string tensor (whose single string has
    # length mnist.IMAGE_PIXELS) to a uint8 tensor with shape
    # [mnist.IMAGE_PIXELS].
    image = tf.decode_raw(features["image_raw"], tf.uint8)
    image.set_shape((mnist.IMAGE_PIXELS))

    # Convert label from a scalar uint8 tensor to an int32 scalar.
    label = tf.cast(features["label"], tf.int32)

    return image, label
//  [python]
def add_loss_op(self, q, target_q):
        """
        Sets the loss of a batch, self.loss is a scalar

        Args:
            q: (tf tensor) shape = (batch_size, num_actions)
            target_q: (tf tensor) shape = (batch_size, num_actions)
        """
        num_actions = self.env.action_space.n

        not_done_mask = tf.abs(tf.cast(self.done_mask, tf.float32) - 1)
        q_samp = self.r + self.config.gamma * not_done_mask * tf.reduce_max(target_q, axis=1)
        a_indices = tf.one_hot(self.a, num_actions)
        q_sa = tf.reduce_sum(q * a_indices, axis=1)
        self.loss = tf.reduce_mean(tf.square(q_samp - q_sa))
//  [python]
def _convert_tf_boxes_to_xyxy(bbox: tf.Tensor, image_size: Sequence[int]):
  """Convert yxyx [0, 1] normalized boxes to xyxy unnormalized format."""
  y0, x0, y1, x1 = tf.split(bbox, 4, axis=-1)
  h = tf.to_float(image_size[0])
  w = tf.to_float(image_size[1])

  y0 = tf.clip_by_value(y0 * h, 0.0, h)
  x0 = tf.clip_by_value(x0 * w, 0.0, w)
  y1 = tf.clip_by_value(y1 * h, 0.0, h)
  x1 = tf.clip_by_value(x1 * w, 0.0, w)

  bbox = tf.concat([x0, y0, x1, y1], axis=-1)
  return bbox
//  [python]
def _convert_tf_boxes_to_xyxy(bbox: tf.Tensor, image_size: Sequence[int]):
  """Convert yxyx [0, 1] normalized boxes to xyxy unnormalized format."""
  y0, x0, y1, x1 = tf.split(bbox, 4, axis=-1)
  h = tf.cast(image_size[0], tf.float32)
  w = tf.cast(image_size[1], tf.float32)

  y0 = tf.clip_by_value(y0 * h, 0.0, h)
  x0 = tf.clip_by_value(x0 * w, 0.0, w)
  y1 = tf.clip_by_value(y1 * h, 0.0, h)
  x1 = tf.clip_by_value(x1 * w, 0.0, w)

  bbox = tf.concat([x0, y0, x1, y1], axis=-1)
  return bbox
//  [python]
def random_apply(func, p, x):
  """Randomly apply function func to x with probability p."""
  return tf.cond(
      tf.less(
          tf.random.uniform([], minval=0, maxval=1, dtype=tf.float32),
          tf.cast(p, tf.float32)), lambda: func(x), lambda: x)
//  [python]
def _multilabel_to_multihot(labels: tf.Tensor, num_classes: int) -> tf.Tensor:
  """Converts labels from multi-label to multi-hot representation.

  In the multi-label representation, labels have shape [...,
  max_num_labels_per_instance] and are integers in [0, num_classes], with -1 for
  padding.

  In the multi-hot representation, labels have shape [..., num_classes + 1] (+1
  due to padding) and are binary vectors with a 1 for each class that applies to
  that instance.

  Args:
    labels: [..., max_num_labels_per_instance] multi-labels.
    num_classes: Number of classes, including padding.

  Returns:
    [..., num_classes + 1] multi-hot label array.
  """

  # Labels are zero-indexed, with -1 for padding (TFDS convention). Add 1 such
  # that multi-hot index 0 means padding:
  labels = tf.one_hot(labels + 1, num_classes)
  labels = tf.reduce_max(labels, axis=-2)  # Combine up multi-labels.

  # Update padding label such that it is 1 iff there are no real labels. This
  # is necessary because multi-labels are padded, which means that the padding
  # label is initally hot for almost all real instances:
  is_padding = tf.to_float(
      tf.reduce_all(labels[..., 1:] == 0, axis=-1, keepdims=True))
  return tf.concat([is_padding, labels[..., 1:]], axis=-1)
//  [python]
def _multilabel_to_multihot(labels: tf.Tensor, num_classes: int) -> tf.Tensor:
  """Converts labels from multi-label to multi-hot representation.

  In the multi-label representation, labels have shape [...,
  max_num_labels_per_instance] and are integers in [0, num_classes], with -1 for
  padding.

  In the multi-hot representation, labels have shape [..., num_classes + 1] (+1
  due to padding) and are binary vectors with a 1 for each class that applies to
  that instance.

  Args:
    labels: [..., max_num_labels_per_instance] multi-labels.
    num_classes: Number of classes, including padding.

  Returns:
    [..., num_classes + 1] multi-hot label array.
  """

  # Labels are zero-indexed, with -1 for padding (TFDS convention). Add 1 such
  # that multi-hot index 0 means padding:
  labels = tf.one_hot(labels + 1, num_classes)
  labels = tf.reduce_max(labels, axis=-2)  # Combine up multi-labels.

  # Update padding label such that it is 1 iff there are no real labels. This
  # is necessary because multi-labels are padded, which means that the padding
  # label is initally hot for almost all real instances:
  is_padding = tf.cast(
      tf.reduce_all(labels[..., 1:] == 0, axis=-1, keepdims=True), labels.dtype)
  return tf.concat([is_padding, labels[..., 1:]], axis=-1)
//  [python]
def apply(self, features: Features) -> Features:
    image = tf.image.convert_image_dtype(features['image'], dtype=tf.float32)

    if self.input_range is not None:
      image = image * (self.input_range[1] -
                       self.input_range[0]) + self.input_range[0]

    image_size = tf.to_int32(tf.shape(image)[:2])

    target = {
        'boxes':
            _convert_tf_boxes_to_xyxy(features[modalities.BOXES], image_size),
        'labels':
            _multilabel_to_multihot(
                labels=features[PER_EXAMPLE_INSTANCE_MULTI_LABELS],
                num_classes=tf.shape(features[modalities.TEXT_QUERIES])[0]),
    }

    out_features = {
        'inputs': image,
        'label': target,
        'queries': features[modalities.TEXT_QUERIES_TOKENIZED],
        'batch_mask': tf.ones((), dtype=tf.float32),
    }

    return detr_transforms.NormalizeBoxes()(out_features)
//  [python]
def apply(self, features: Features) -> Features:
    image = tf.image.convert_image_dtype(features['image'], dtype=tf.float32)

    if self.input_range is not None:
      image = image * (self.input_range[1] -
                       self.input_range[0]) + self.input_range[0]

    image_size = tf.cast(tf.shape(image)[:2], dtype=tf.int32)

    target = {
        'boxes':
            _convert_tf_boxes_to_xyxy(features[modalities.BOXES], image_size),
        'labels':
            _multilabel_to_multihot(
                labels=features[PER_EXAMPLE_INSTANCE_MULTI_LABELS],
                num_classes=tf.shape(features[modalities.TEXT_QUERIES])[0]),
    }

    out_features = {
        'inputs': image,
        'label': target,
        'queries': features[modalities.TEXT_QUERIES_TOKENIZED],
        'batch_mask': tf.ones((), dtype=tf.float32),
    }

    return detr_transforms.NormalizeBoxes()(out_features)
//  [python]
def __call__(self, features: Features) -> Features:
    image = tf.cast(features['image'], tf.float32) / 255.0

    boxes, text_labels = [], []
    if self.include_objects:
      boxes.append(features['objects']['bbox'])  # float32, in range [0, 1].
      text_labels.append(features['objects']['name'])

    if self.include_regions:
      boxes.append(features['regions']['bbox'])  # float32, in range [0, 1].
      text_labels.append(features['regions']['phrase'])

    # Combined objects and regions.
    if boxes:
      boxes = tf.concat(boxes, axis=0)
      text_labels = tf.concat(text_labels, axis=0)
    else:
      raise ValueError('Either objects or regions should be included in VG.')

    # Remove empty text labels.
    # pylint: disable=g-explicit-bool-comparison
    boxes = boxes[text_labels != PADDING_QUERY]
    text_labels = text_labels[text_labels != PADDING_QUERY]
    # pylint: enable=g-explicit-bool-comparison

    # Visual Genome performs better without prompting:
    # First, remove any markers that might be present in the labels:
    text_labels = remove_promptability_marker(text_labels)
    if not self.is_promptable:
      text_labels = mark_not_promptable(text_labels)

    features_new = {
        modalities.IMAGE:
            image,
        modalities.BOXES:
            boxes,
        modalities.INSTANCE_TEXT_LABELS:
            text_labels,
        modalities.NEGATIVE_LABELS:
            tf.fill([1], -1),  # Dummy padding.
        modalities.NEGATIVE_TEXT_LABELS:
            tf.fill([1], PADDING_QUERY),  # Dummy padding.
        # There are no labels so everything just 0. Don't set to -1 as mosaic
        # operation will filter padding.
        modalities.INSTANCE_LABELS:
            tf.zeros_like(text_labels, dtype=tf.int32),
        modalities.CROWD:
            tf.zeros_like(text_labels, dtype=tf.int32),
    }

    if 'rng' in features:
      features_new[image_ops.SEED_KEY] = features['rng']
    return features_new
//  [python]
def df_to_ds(X, y, shuffle=False):
  ds = tf.data.Dataset.from_tensor_slices((
    X.values,
    tf.one_hot(tf.to_int32(y.values), depth=2, axis=-1)
  ))
  
  if shuffle:
    ds = ds.shuffle(1000, seed=1, reshuffle_each_iteration=True)
  
  return ds.map(lambda text, label: ({ tweet_text_feature: text }, label)).batch(params['batch_size'])
//  [python]
def df_to_ds(X, y, shuffle=False):
  ds = tf.data.Dataset.from_tensor_slices((
    X.values,
    tf.one_hot(tf.cast(y.values, tf.int32), depth=2, axis=-1)
  ))
  
  if shuffle:
    ds = ds.shuffle(1000, seed=1, reshuffle_each_iteration=True)
  
  return ds.map(lambda text, label: ({ tweet_text_feature: text }, label)).batch(params['batch_size'])
//  [python]
def test_saved_model_creation_extended(self):
        config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
        config.output_hidden_states = True
        config.output_attentions = True

        if hasattr(config, "use_cache"):
            config.use_cache = True

        encoder_seq_length = getattr(self.model_tester, "encoder_seq_length", self.model_tester.seq_length)
        encoder_key_length = getattr(self.model_tester, "key_length", encoder_seq_length)

        for model_class in self.all_model_classes[:2]:
            class_inputs_dict = self._prepare_for_class(inputs_dict, model_class)
            model = model_class(config)
            model.build_in_name_scope()
            num_out = len(model(class_inputs_dict))

            for key in list(class_inputs_dict.keys()):
                # Remove keys not in the serving signature, as the SavedModel will not be compiled to deal with them
                if key not in model.input_signature:
                    del class_inputs_dict[key]
                # Check it's a tensor, in case the inputs dict has some bools in it too
                elif isinstance(class_inputs_dict[key], tf.Tensor) and class_inputs_dict[key].dtype.is_integer:
                    class_inputs_dict[key] = tf.cast(class_inputs_dict[key], tf.int32)

            if set(class_inputs_dict.keys()) != set(model.input_signature.keys()):
                continue  # Some models have inputs that the preparation functions don't create, we skip those

            with tempfile.TemporaryDirectory() as tmpdirname:
                model.save_pretrained(tmpdirname, saved_model=True)
                saved_model_dir = os.path.join(tmpdirname, "saved_model", "1")
                model = keras.models.load_model(saved_model_dir)
                outputs = model(class_inputs_dict)

                if self.is_encoder_decoder:
                    output_hidden_states = outputs["encoder_hidden_states"]
                    output_attentions = outputs["encoder_attentions"]
                else:
                    output_hidden_states = outputs["hidden_states"]
                    output_attentions = outputs["attentions"]

                self.assertEqual(len(outputs), num_out)

                expected_num_layers = getattr(
                    self.model_tester, "expected_num_hidden_layers", self.model_tester.num_hidden_layers + 1
                )

                self.assertEqual(len(output_hidden_states), expected_num_layers)
                self.assertListEqual(
                    list(output_hidden_states[0].shape[-2:]),
                    [self.model_tester.seq_length, self.model_tester.hidden_size],
                )

                self.assertEqual(len(output_attentions), self.model_tester.num_hidden_layers)
                self.assertListEqual(
                    list(output_attentions[0].shape[-3:]),
                    [self.model_tester.num_attention_heads, encoder_seq_length, encoder_key_length],
                )
//  [python]
def call(self, inputs, mode="training"):  # pytype: disable=signature-mismatch  # overriding-default-value-checks
    input_shape = tf_utils.get_shape_list(inputs["input_ids"], expected_rank=3)
    batch_size, num_docs, len_passage = (input_shape[0], input_shape[1],
                                         input_shape[2])
    input_ids = tf.reshape(inputs["input_ids"], [-1, len_passage])
    input_mask = tf.reshape(inputs["input_mask"], [-1, len_passage])
    segment_ids = tf.reshape(inputs["segment_ids"], [-1, len_passage])
    all_encoder_outputs, _ = self.bert_layer(
        [input_ids, input_mask, segment_ids])
    encoder_outputs = tf.reshape(
        all_encoder_outputs[-1],
        [batch_size, num_docs, len_passage, self.params.hidden_size])
    doc_attention_mask = tf.reshape(
        tf.to_int32(
            tf.math.count_nonzero(input_mask, axis=1, dtype=tf.int32) > 2), [batch_size, num_docs])

    doc_attention_probs = self.doc_attention(encoder_outputs,
                                             doc_attention_mask)
    encoder_decoder_attention_bias = decoder.get_attention_bias(
        inputs["input_ids"],
        bias_type="multi_cross",
        padding_value=self.params.pad_token_id)

    if mode == "train":
      target_length = tf_utils.get_shape_list(
          inputs["target_ids"], expected_rank=2)[1]
      doc_attention_probs = self._expand_doc_attention_probs(
          doc_attention_probs, target_length)
      self_attention_bias = decoder.get_attention_bias(
          inputs["target_ids"], bias_type="decoder_self")
      decoder_inputs = dict(
          attention_bias=encoder_decoder_attention_bias,
          self_attention_bias=self_attention_bias,
          target_ids=inputs["target_ids"],
          all_encoder_outputs=encoder_outputs,
          doc_attention_probs=doc_attention_probs)
      decoder_outputs = self.decoder_layer(decoder_inputs)
      return self.train_decode(decoder_outputs)

    # Adds encoder output and attention bias to the cache.
    if self.params.use_cache:
      cache = self._init_cache(batch_size)
    else:
      cache = {}
    cache["all_encoder_outputs"] = [encoder_outputs]
    cache["attention_bias"] = encoder_decoder_attention_bias
    cache["doc_attention_probs"] = doc_attention_probs

    start_token_ids = tf.ones([batch_size],
                              tf.int32) * self.params.start_token_id
    decoded_ids, scores = self.predict_decode(start_token_ids, cache)
    if mode == "predict":
      return decoded_ids[:, :self.params.beam_size,
                         1:], scores[:, :self.params.beam_size]

    top_decoded_ids = decoded_ids[:, 0, 1:]
    target_length = tf_utils.get_shape_list(top_decoded_ids)[-1]
    decoder_inputs = dict(
        attention_bias=encoder_decoder_attention_bias,
        all_encoder_outputs=[encoder_outputs],
        doc_attention_probs=self._expand_doc_attention_probs(
            doc_attention_probs, target_length))
    return self._get_logits_for_decode_ids(decoder_inputs, top_decoded_ids)
//  [python]
def call(self, inputs, mode="training"):  # pytype: disable=signature-mismatch  # overriding-default-value-checks
    input_shape = tf_utils.get_shape_list(inputs["input_ids"], expected_rank=3)
    batch_size, num_docs, len_passage = (input_shape[0], input_shape[1],
                                         input_shape[2])
    input_ids = tf.reshape(inputs["input_ids"], [-1, len_passage])
    input_mask = tf.reshape(inputs["input_mask"], [-1, len_passage])
    segment_ids = tf.reshape(inputs["segment_ids"], [-1, len_passage])
    all_encoder_outputs, _ = self.bert_layer(
        [input_ids, input_mask, segment_ids])
    encoder_outputs = tf.reshape(
        all_encoder_outputs[-1],
        [batch_size, num_docs, len_passage, self.params.hidden_size])
    doc_attention_mask = tf.reshape(
        tf.cast(
            tf.math.count_nonzero(input_mask, axis=1, dtype=tf.int32) > 2,
            tf.int32), [batch_size, num_docs])

    doc_attention_probs = self.doc_attention(encoder_outputs,
                                             doc_attention_mask)
    encoder_decoder_attention_bias = decoder.get_attention_bias(
        inputs["input_ids"],
        bias_type="multi_cross",
        padding_value=self.params.pad_token_id)

    if mode == "train":
      target_length = tf_utils.get_shape_list(
          inputs["target_ids"], expected_rank=2)[1]
      doc_attention_probs = self._expand_doc_attention_probs(
          doc_attention_probs, target_length)
      self_attention_bias = decoder.get_attention_bias(
          inputs["target_ids"], bias_type="decoder_self")
      decoder_inputs = dict(
          attention_bias=encoder_decoder_attention_bias,
          self_attention_bias=self_attention_bias,
          target_ids=inputs["target_ids"],
          all_encoder_outputs=encoder_outputs,
          doc_attention_probs=doc_attention_probs)
      decoder_outputs = self.decoder_layer(decoder_inputs)
      return self.train_decode(decoder_outputs)

    # Adds encoder output and attention bias to the cache.
    if self.params.use_cache:
      cache = self._init_cache(batch_size)
    else:
      cache = {}
    cache["all_encoder_outputs"] = [encoder_outputs]
    cache["attention_bias"] = encoder_decoder_attention_bias
    cache["doc_attention_probs"] = doc_attention_probs

    start_token_ids = tf.ones([batch_size],
                              tf.int32) * self.params.start_token_id
    decoded_ids, scores = self.predict_decode(start_token_ids, cache)
    if mode == "predict":
      return decoded_ids[:, :self.params.beam_size,
                         1:], scores[:, :self.params.beam_size]

    top_decoded_ids = decoded_ids[:, 0, 1:]
    target_length = tf_utils.get_shape_list(top_decoded_ids)[-1]
    decoder_inputs = dict(
        attention_bias=encoder_decoder_attention_bias,
        all_encoder_outputs=[encoder_outputs],
        doc_attention_probs=self._expand_doc_attention_probs(
            doc_attention_probs, target_length))
    return self._get_logits_for_decode_ids(decoder_inputs, top_decoded_ids)
//  [python]
def _train(self, experience, weights=None):
    """Updates the policy based on the data in `experience`.

    Note that `experience` should only contain data points that this agent has
    not previously seen. If `experience` comes from a replay buffer, this buffer
    should be cleared between each call to `train`.

    Args:
      experience: A batch of experience data in the form of a `Trajectory`.
      weights: Unused.

    Returns:
      A `LossInfo` containing the loss *before* the training step is taken.
        Note that the loss does not depend on policy state and comes directly
        from the experience (and is therefore not differentiable).

        In most cases, if `weights` is provided, the entries of this tuple will
        have been calculated with the weights.  Note that each Agent chooses
        its own method of applying weights.
    """
    del weights  # unused
    experience = self._as_trajectory(experience)
    reward = experience.reward
    log_prob = policy_step.get_log_probability(experience.policy_info)
    action = experience.action
    update_value = exp3_update_value(reward, log_prob)
    weight_update = selective_sum(
        values=update_value, partitions=action, num_partitions=self.num_actions
    )
    tf.compat.v1.assign_add(self._weights, weight_update)

    batch_size = tf.size(reward)
    self._train_step_counter.assign_add(batch_size)

    return tf_agent.LossInfo(loss=-tf.reduce_sum(experience.reward), extra=())
//  [python]
def _train(self, experience, weights=None):
    """Updates the policy based on the data in `experience`.

    Note that `experience` should only contain data points that this agent has
    not previously seen. If `experience` comes from a replay buffer, this buffer
    should be cleared between each call to `train`.

    Args:
      experience: A batch of experience data in the form of a `Trajectory`.
      weights: Unused.

    Returns:
      A `LossInfo` containing the loss *before* the training step is taken.
        Note that the loss does not depend on policy state and comes directly
        from the experience (and is therefore not differentiable).

        In most cases, if `weights` is provided, the entries of this tuple will
        have been calculated with the weights.  Note that each Agent chooses
        its own method of applying weights.
    """
    del weights  # unused
    experience = self._as_trajectory(experience)
    reward = experience.reward
    log_prob = policy_step.get_log_probability(experience.policy_info)
    action = experience.action
    update_value = exp3_update_value(reward, log_prob)
    weight_update = selective_sum(
        values=update_value, partitions=action, num_partitions=self.num_actions
    )
    tf.compat.v1.assign_add(self._weights, weight_update)

    batch_size = tf.cast(tf.size(reward), dtype=tf.int64)
    self._train_step_counter.assign_add(batch_size)

    return tf_agent.LossInfo(loss=-tf.reduce_sum(experience.reward), extra=())
//  [python]
def add_evaluation_step(result_tensor, ground_truth_tensor):
    """Inserts the operations we need to evaluate the accuracy of our results.

  Args:
    result_tensor: The new final node that produces results.
    ground_truth_tensor: The node we feed ground truth data
    into.

  Returns:
    Tuple of (evaluation step, prediction).
  """
    with tf.name_scope('accuracy'):
        with tf.name_scope('correct_prediction'):
            prediction = tf.argmax(result_tensor, 1)
            correct_prediction = tf.equal(prediction, ground_truth_tensor)
        with tf.name_scope('accuracy'):
            evaluation_step = tf.reduce_mean(
                tf.to_float(correct_prediction))
    tf.summary.scalar('accuracy', evaluation_step)
    return evaluation_step, prediction
//  [python]
def add_evaluation_step(result_tensor, ground_truth_tensor):
    """Inserts the operations we need to evaluate the accuracy of our results.

  Args:
    result_tensor: The new final node that produces results.
    ground_truth_tensor: The node we feed ground truth data
    into.

  Returns:
    Tuple of (evaluation step, prediction).
  """
    with tf.name_scope('accuracy'):
        with tf.name_scope('correct_prediction'):
            prediction = tf.argmax(result_tensor, 1)
            correct_prediction = tf.equal(prediction, ground_truth_tensor)
        with tf.name_scope('accuracy'):
            evaluation_step = tf.reduce_mean(
                tf.cast(correct_prediction, tf.float32))
    tf.summary.scalar('accuracy', evaluation_step)
    return evaluation_step, prediction
//  [python]
def gaussian_blur(image, kernel_size, sigma, padding='SAME'):
  """Blurs the given image with separable convolution.


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

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


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

  Returns:
    A Tensor representing the blurred image.
  """
  radius = tf.cast(kernel_size / 2, dtype=tf.int32)
  kernel_size = radius * 2 + 1
  x = tf.cast(tf.range(-radius, radius + 1), dtype=tf.float32)
  blur_filter = tf.exp(-tf.pow(x, 2.0) /
                       (2.0 * tf.pow(tf.cast(sigma, dtype=tf.float32), 2.0)))
  blur_filter /= tf.reduce_sum(blur_filter)
  # One vertical and one horizontal filter.
  blur_v = tf.reshape(blur_filter, [kernel_size, 1, 1, 1])
  blur_h = tf.reshape(blur_filter, [1, kernel_size, 1, 1])
  num_channels = tf.shape(image)[-1]
  blur_h = tf.tile(blur_h, [1, 1, num_channels, 1])
  blur_v = tf.tile(blur_v, [1, 1, num_channels, 1])
  expand_batch_dim = image.shape.ndims == 3
  if expand_batch_dim:
    # Tensorflow requires batched input to convolutions, which we can fake with
    # an extra dimension.
    image = tf.expand_dims(image, axis=0)
  blurred = tf.nn.depthwise_conv2d(
      image, blur_h, strides=[1, 1, 1, 1], padding=padding)
  blurred = tf.nn.depthwise_conv2d(
      blurred, blur_v, strides=[1, 1, 1, 1], padding=padding)
  if expand_batch_dim:
    blurred = tf.squeeze(blurred, axis=0)
  return blurred
//  [python]
def add_jpeg_decoding(module_spec):
    """Adds operations that perform JPEG decoding and resizing to the graph..

  Args:
    module_spec: The hub.ModuleSpec for the image module being used.

  Returns:
    Tensors for the node to feed JPEG data into, and the output of the
      preprocessing steps.
  """
    input_height, input_width = hub.get_expected_image_size(module_spec)
    input_depth = hub.get_num_image_channels(module_spec)
    jpeg_data = tf.compat.v1.placeholder(tf.string, name='DecodeJPGInput')
    decoded_image = tf.image.decode_jpeg(jpeg_data, channels=input_depth)
    # Convert from full range of uint8 to range [0,1] of float32.
    decoded_image_as_float = tf.image.convert_image_dtype(
        decoded_image, tf.float32)
    decoded_image_4d = tf.expand_dims(decoded_image_as_float, 0)
    resize_shape = tf.stack([input_height, input_width])
    resize_shape_as_int = tf.cast(resize_shape, dtype=tf.int32)
    resized_image = tf.image.resize_bilinear(decoded_image_4d,
                                             resize_shape_as_int)
    return jpeg_data, resized_image
//  [python]
def call(self, inputs, mask=None):
    if mask is None:
      raise ValueError('Inputs to `MaskedAverage` need to be masked.')

    # Get inputs with masked values replaced by 0's.
    mask = tf.expand_dims(tf.cast(mask, 'float32'), axis=-1)
    masked_input = inputs * mask

    # Perform a mask-weighted sum across the axis, divide by the weights to
    # get the weighted average.
    weighted_summed_input = tf.reduce_sum(masked_input, axis=self.axis)
    total_weights = tf.reduce_sum(mask, axis=self.axis)
    return tf.math.divide_no_nan(weighted_summed_input, total_weights)
//  [python]
def correct_rate(labels, predictions, weights):
    weights = tf.ones_like(weights, tf.int32)

    labels = tf.to_int32(labels)
    # with tf.Session() as sess:
    #     sess.run(tf.tables_initializer())
    #     result = sess.run((labels, weights))
    #     print(result)

    label_seq = tf.multiply(labels, weights, name='labels')

    predictions = tf.to_int32(predictions)
    prediction_seq = tf.multiply(predictions, weights, name='predictions')

    flag = tf.math.equal(label_seq, prediction_seq)

    correct = tf.math.reduce_all(flag, -1)
    
    fake_predictions = tf.to_int32(correct, name='fake_prediction')

    fake_golden = tf.ones(tf.shape(fake_predictions), dtype=tf.int32, name='fake_golden')

    # print(fake_golden, fake_predictions)

    # return tf.metrics.accuracy(fake_golden, fake_predictions), fake_golden, fake_predictions
    return tf.metrics.accuracy(fake_golden, fake_predictions)
//  [python]
def correct_rate(labels, predictions, weights):
    weights = tf.ones_like(weights, tf.int32)

    labels = tf.cast(labels, tf.int32)
    # with tf.Session() as sess:
    #     sess.run(tf.tables_initializer())
    #     result = sess.run((labels, weights))
    #     print(result)

    label_seq = tf.multiply(labels, weights, name='labels')

    predictions = tf.cast(predictions, tf.int32)
    prediction_seq = tf.multiply(predictions, weights, name='predictions')

    flag = tf.math.equal(label_seq, prediction_seq)

    correct = tf.math.reduce_all(flag, -1)
    
    fake_predictions = tf.to_int32(correct, name='fake_prediction')

    fake_golden = tf.ones(tf.shape(fake_predictions), dtype=tf.int32, name='fake_golden')

    # print(fake_golden, fake_predictions)

    # return tf.metrics.accuracy(fake_golden, fake_predictions), fake_golden, fake_predictions
    return tf.metrics.accuracy(fake_golden, fake_predictions)