jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def build(self, images):
    """Builds a ResNet50 embedder for the input images.

    It assumes that the range of the pixel values in the images tensor is
      [0,255] and should be castable to tf.uint8.

    Args:
      images: a tensor that contains the input images which has the shape of
          NxTxHxWx3 where N is the batch size, T is the maximum length of the
          sequence, H and W are the height and width of the images and C is the
          number of channels.
    Returns:
      The embedding of the input image with the shape of NxTxL where L is the
        embedding size of the output.

    Raises:
      ValueError: if the shape of the input does not agree with the expected
      shape explained in the Args section.
    """
    shape = images.get_shape().as_list()
    if len(shape) != 5:
      raise ValueError(
          'The tensor shape should have 5 elements, {} is provided'.format(
              len(shape)))
    if shape[4] != 3:
      raise ValueError('Three channels are expected for the input image')

    images = tf.cast(images, tf.uint8)
    images = tf.reshape(images,
                        [shape[0] * shape[1], shape[2], shape[3], shape[4]])
    with slim.arg_scope(resnet_v2.resnet_arg_scope()):

      def preprocess_fn(x):
        x = tf.expand_dims(x, 0)
        x = tf.image.resize_bilinear(x, [299, 299],
                                       align_corners=False)
        return(tf.squeeze(x, [0]))

      images = tf.map_fn(preprocess_fn, images, dtype=tf.float32)

      net, _ = resnet_v2.resnet_v2_50(
          images, is_training=False, global_pool=True)
      output = tf.reshape(net, [shape[0], shape[1], -1])
      return output
//  [python]
def build_tensors_for_augmentation(
    frame: Dict[str, tf.Tensor]
) -> Tuple[tf.Tensor, tf.Tensor]:
    """Builds tensors for data augmentation from an input frame.

    Args:
      frame: a dictionary of feature tensors from a Waymo Open Dataset frame

    Returns:
      A dictionary of two tensors with keys "point_clouds" and "bounding_boxes"
      and values which are tensors of shapes [num points, num features] and
      [num boxes, num features]).
    """
    assert_waymo_open_dataset_installed(
        "keras_cv.datasets.waymo.build_tensors_for_augmentation()"
    )
    point_cloud = tf.concat(
        [
            frame["point_xyz"][tf.newaxis, ...],
            frame["point_feature"][tf.newaxis, ...],
            tf.to_float(frame["point_mask"])[tf.newaxis, :, tf.newaxis],
        ],
        axis=-1,
    )
    boxes = tf.concat(
        [
            frame["label_box"][tf.newaxis, :],
            tf.to_float(frame["label_box_class"])[
                tf.newaxis, :, tf.newaxis
            ],
            tf.to_float(frame["label_box_mask"])[
                tf.newaxis, :, tf.newaxis
            ],
            tf.to_float(frame["label_box_density"])[
                tf.newaxis, :, tf.newaxis
            ],
            tf.to_float(frame["label_box_detection_difficulty"])[
                tf.newaxis, :, tf.newaxis
            ],
        ],
        axis=-1,
    )
    return {
        "point_clouds": tf.squeeze(point_cloud, axis=0),
        "bounding_boxes": tf.squeeze(boxes, axis=0),
    }
//  [python]
def build_tensors_for_augmentation(
    frame: Dict[str, tf.Tensor]
) -> Tuple[tf.Tensor, tf.Tensor]:
    """Builds tensors for data augmentation from an input frame.

    Args:
      frame: a dictionary of feature tensors from a Waymo Open Dataset frame

    Returns:
      A dictionary of two tensors with keys "point_clouds" and "bounding_boxes"
      and values which are tensors of shapes [num points, num features] and
      [num boxes, num features]).
    """
    assert_waymo_open_dataset_installed(
        "keras_cv.datasets.waymo.build_tensors_for_augmentation()"
    )
    point_cloud = tf.concat(
        [
            frame["point_xyz"][tf.newaxis, ...],
            frame["point_feature"][tf.newaxis, ...],
            tf.cast(frame["point_mask"], tf.float32)[tf.newaxis, :, tf.newaxis],
        ],
        axis=-1,
    )
    boxes = tf.concat(
        [
            frame["label_box"][tf.newaxis, :],
            tf.cast(frame["label_box_class"], tf.float32)[
                tf.newaxis, :, tf.newaxis
            ],
            tf.cast(frame["label_box_mask"], tf.float32)[
                tf.newaxis, :, tf.newaxis
            ],
            tf.cast(frame["label_box_density"], tf.float32)[
                tf.newaxis, :, tf.newaxis
            ],
            tf.cast(frame["label_box_detection_difficulty"], tf.float32)[
                tf.newaxis, :, tf.newaxis
            ],
        ],
        axis=-1,
    )
    return {
        "point_clouds": tf.squeeze(point_cloud, axis=0),
        "bounding_boxes": tf.squeeze(boxes, axis=0),
    }
//  [python]
def roi_pooling(feature_maps, rois, img_shape):
  '''
  Here use roi warping as roi_pooling
  :param featuremaps_dict: feature map to crop
  :param rois: shape is [-1, 4]. [x1, y1, x2, y2]
  :return:
  '''
  with tf.variable_scope('ROI_Warping'):
    img_h, img_w = tf.to_float(img_shape[1]), tf.to_float(img_shape[2])
    N = tf.shape(rois)[0]
    x1, y1, x2, y2 = tf.unstack(rois, axis=1)

    normalized_x1 = x1 / img_w
    normalized_x2 = x2 / img_w
    normalized_y1 = y1 / img_h
    normalized_y2 = y2 / img_h

    normalized_rois = tf.transpose(
        tf.stack([normalized_y1, normalized_x1, normalized_y2, normalized_x2]), name='get_normalized_rois')

    normalized_rois = tf.stop_gradient(normalized_rois)

    cropped_roi_features = tf.image.crop_and_resize(feature_maps, normalized_rois,
                                                    box_ind=tf.zeros(shape=[N, ],
                                                                     dtype=tf.int32),
                                                    crop_size=[cfgs.ROI_SIZE, cfgs.ROI_SIZE],
                                                    name='CROP_AND_RESIZE'
                                                    )
    roi_features = slim.max_pool2d(cropped_roi_features,
                                  [cfgs.ROI_POOL_KERNEL_SIZE, cfgs.ROI_POOL_KERNEL_SIZE],
                                  stride=cfgs.ROI_POOL_KERNEL_SIZE)

  return roi_features
//  [python]
def roi_pooling(feature_maps, rois, img_shape):
  '''
  Here use roi warping as roi_pooling
  :param featuremaps_dict: feature map to crop
  :param rois: shape is [-1, 4]. [x1, y1, x2, y2]
  :return:
  '''
  with tf.variable_scope('ROI_Warping'):
    img_h, img_w = tf.cast(img_shape[1], tf.float32), tf.cast(img_shape[2], tf.float32)
    N = tf.shape(rois)[0]
    x1, y1, x2, y2 = tf.unstack(rois, axis=1)

    normalized_x1 = x1 / img_w
    normalized_x2 = x2 / img_w
    normalized_y1 = y1 / img_h
    normalized_y2 = y2 / img_h

    normalized_rois = tf.transpose(
        tf.stack([normalized_y1, normalized_x1, normalized_y2, normalized_x2]), name='get_normalized_rois')

    normalized_rois = tf.stop_gradient(normalized_rois)

    cropped_roi_features = tf.image.crop_and_resize(feature_maps, normalized_rois,
                                                    box_ind=tf.zeros(shape=[N, ],
                                                                     dtype=tf.int32),
                                                    crop_size=[cfgs.ROI_SIZE, cfgs.ROI_SIZE],
                                                    name='CROP_AND_RESIZE'
                                                    )
    roi_features = slim.max_pool2d(cropped_roi_features,
                                  [cfgs.ROI_POOL_KERNEL_SIZE, cfgs.ROI_POOL_KERNEL_SIZE],
                                  stride=cfgs.ROI_POOL_KERNEL_SIZE)

  return roi_features
//  [python]
def train(self):
        y = tf.one_hot(self.y, args.class_size)
        loss = tf.nn.softmax_cross_entropy_with_logits(labels=y, logits=self.logits)
        self.loss = tf.reduce_mean(loss)
        self.train_op = tf.train.AdamOptimizer(args.learning_rate).minimize(self.loss)
        # batch_normalization计算均值和方差
        update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
        self.train_op = tf.group([self.train_op, update_ops])
        self.predict = tf.argmax(self.logits, axis=1, name="predictions")
        correct_prediction = tf.equal(tf.cast(self.predict, tf.int32), self.y)
        self.accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32), name="Accuracy")
//  [python]
def read_smpl_data_temporal(filename_queue):
    """
    Parses smpl temporal Example proto.
    """
    with tf.name_scope(None, 'read_smpl_data_temporal', [filename_queue]):
        reader = tf.TFRecordReader()
        _, example_serialized = reader.read(filename_queue)

        feature_map = {
            'delta_pose': tf.FixedLenFeature((72, ), dtype=tf.float32),
            'pose': tf.FixedLenFeature((72, ), dtype=tf.float32),
        }

        features = tf.parse_single_example(example_serialized, feature_map)
        pose = tf.to_float(features['pose'])
        delta_pose = tf.to_float(features['delta_pose'])
        return pose, delta_pose
//  [python]
def read_smpl_data_temporal(filename_queue):
    """
    Parses smpl temporal Example proto.
    """
    with tf.name_scope(None, 'read_smpl_data_temporal', [filename_queue]):
        reader = tf.TFRecordReader()
        _, example_serialized = reader.read(filename_queue)

        feature_map = {
            'delta_pose': tf.FixedLenFeature((72, ), dtype=tf.float32),
            'pose': tf.FixedLenFeature((72, ), dtype=tf.float32),
        }

        features = tf.parse_single_example(example_serialized, feature_map)
        pose = tf.cast(features['pose'], dtype=tf.float32)
        delta_pose = tf.cast(features['delta_pose'], dtype=tf.float32)
        return pose, delta_pose
//  [python]
def classification_loss(truth, predicted, weights=None, is_one_hot=True):
  """A cross entropy loss.

  Computes the mean of cross entropy losses for all pairs of true labels and
  predictions. It wraps around a tf implementation of the cross entropy loss
  with additional reformating of the inputs. If the truth and predicted are
  n-rank Tensors with n > 2, then these are reshaped to 2-rank Tensors. It
  allows for truth to be specified as one hot vector or class indices. Finally,
  a weight can be specified for each element in truth and predicted.

  Args:
    truth: an n-rank or (n-1)-rank Tensor containing labels. If is_one_hot is
      True, then n-rank Tensor is expected, otherwise (n-1) rank one.
    predicted: an n-rank float Tensor containing prediction probabilities.
    weights: an (n-1)-rank float Tensor of weights
    is_one_hot: a boolean.

  Returns:
    A TF float scalar.
  """
  num_labels = predicted.get_shape().as_list()[-1]
  if not is_one_hot:
    truth = tf.reshape(truth, [-1])
    truth = tf.one_hot(
        truth, depth=num_labels, on_value=1.0, off_value=0.0, axis=-1)
  else:
    truth = tf.reshape(truth, [-1, num_labels])
  predicted = tf.reshape(predicted, [-1, num_labels])
  losses = tf.nn.softmax_cross_entropy_with_logits(
      labels=truth, logits=predicted)
  if weights is not None:
    losses = tf.boolean_mask(losses,
                             tf.to_bool(tf.reshape(weights, [-1])))
  return tf.reduce_mean(losses)
//  [python]
def classification_loss(truth, predicted, weights=None, is_one_hot=True):
  """A cross entropy loss.

  Computes the mean of cross entropy losses for all pairs of true labels and
  predictions. It wraps around a tf implementation of the cross entropy loss
  with additional reformating of the inputs. If the truth and predicted are
  n-rank Tensors with n > 2, then these are reshaped to 2-rank Tensors. It
  allows for truth to be specified as one hot vector or class indices. Finally,
  a weight can be specified for each element in truth and predicted.

  Args:
    truth: an n-rank or (n-1)-rank Tensor containing labels. If is_one_hot is
      True, then n-rank Tensor is expected, otherwise (n-1) rank one.
    predicted: an n-rank float Tensor containing prediction probabilities.
    weights: an (n-1)-rank float Tensor of weights
    is_one_hot: a boolean.

  Returns:
    A TF float scalar.
  """
  num_labels = predicted.get_shape().as_list()[-1]
  if not is_one_hot:
    truth = tf.reshape(truth, [-1])
    truth = tf.one_hot(
        truth, depth=num_labels, on_value=1.0, off_value=0.0, axis=-1)
  else:
    truth = tf.reshape(truth, [-1, num_labels])
  predicted = tf.reshape(predicted, [-1, num_labels])
  losses = tf.nn.softmax_cross_entropy_with_logits(
      labels=truth, logits=predicted)
  if weights is not None:
    losses = tf.boolean_mask(losses,
                             tf.cast(tf.reshape(weights, [-1]), dtype=tf.bool))
  return tf.reduce_mean(losses)
//  [python]
def process_tensors(data, N, indices, shape, dtype, default=0, name=None):
    """
    Wrapper function for processing sparse tensors outputted by variable length
    features.

    1. Converts sparse tensors to dense tensors.
    2. Reshapes and fills in default value.
    3. Subsamples tensor based on indices.

    Args:
        data (SparseTensor): Sparse Tensor to convert.
        N (scalar Tensor): Total number of frames.
        indices (Tensor): Indices to keep.
        shape (tuple): Shape of tensor.
        dtype (dtype): Data type if need to cast.
        default (float or str): Default value when converting to dense
            tensor or if tensor is empty.
        name (str): scope.

    Returns:
        Tensor.
    """
    with tf.name_scope(name, 'process_tensors', [data]):
        if dtype:
            data = tf.cast(data, dtype)
        dense_tensor = tf.sparse_tensor_to_dense(data, default_value=default)
        dense_tensor = tf.cond(
            tf.equal(tf.size(dense_tensor), 0),
            lambda: tf.cast(tf.fill(shape, default), dtype),
            lambda: tf.reshape(subsample(dense_tensor, N, indices), shape)
        )
    return dense_tensor
//  [python]
def prepare_cifar(x, y):

    x = t.cast(x, tf.float32)
    y = t.cast(y, tf.int32)
    return x, y
//  [python]
def prepare_cifar(x, y):

    x = tf.cast(x, tf.float32)
    y = tf.cast(y, tf.int32)
    return x, y
//  [python]
def _reorg_boxes(self, boxes, info, num_detections):
    """Scale and Clean boxes prior to Evaluation."""
    mask = tf.sequence_mask(num_detections, maxlen=tf.shape(boxes)[1])
    mask = tf.expand_dims(mask, axis=-1)

    # Denormalize the boxes by the shape of the image
    inshape = tf.expand_dims(info[:, 1, :], axis=1)
    ogshape = tf.expand_dims(info[:, 0, :], axis=1)
    scale = tf.expand_dims(info[:, 2, :], axis=1)
    offset = tf.expand_dims(info[:, 3, :], axis=1)

    boxes = box_ops.denormalize_boxes(boxes, inshape)
    boxes = box_ops.clip_boxes(boxes, inshape)
    boxes += tf.tile(offset, [1, 1, 2])
    boxes /= tf.tile(scale, [1, 1, 2])
    boxes = box_ops.clip_boxes(boxes, ogshape)

    # Mask the boxes for usage
    boxes *= mask
    boxes += (mask - 1)
    return boxes
//  [python]
def _reorg_boxes(self, boxes, info, num_detections):
    """Scale and Clean boxes prior to Evaluation."""
    mask = tf.sequence_mask(num_detections, maxlen=tf.shape(boxes)[1])
    mask = tf.cast(tf.expand_dims(mask, axis=-1), boxes.dtype)

    # Denormalize the boxes by the shape of the image
    inshape = tf.expand_dims(info[:, 1, :], axis=1)
    ogshape = tf.expand_dims(info[:, 0, :], axis=1)
    scale = tf.expand_dims(info[:, 2, :], axis=1)
    offset = tf.expand_dims(info[:, 3, :], axis=1)

    boxes = box_ops.denormalize_boxes(boxes, inshape)
    boxes = box_ops.clip_boxes(boxes, inshape)
    boxes += tf.tile(offset, [1, 1, 2])
    boxes /= tf.tile(scale, [1, 1, 2])
    boxes = box_ops.clip_boxes(boxes, ogshape)

    # Mask the boxes for usage
    boxes *= mask
    boxes += (mask - 1)
    return boxes
//  [python]
def train_step(self, inputs, model, optimizer, metrics=None):
    """Train Step.

    Forward step and backwards propagate the model.

    Args:
      inputs: a dictionary of input tensors.
      model: the model, forward pass definition.
      optimizer: the optimizer for this training step.
      metrics: a nested structure of metrics objects.

    Returns:
      A dictionary of logs.
    """
    image, label = inputs

    with tf.GradientTape(persistent=False) as tape:
      # Compute a prediction
      y_pred = model(image, training=True)

      # Cast to float32 for gradietn computation
      y_pred = tf.nest.map_structure(lambda x: tf.to_float(x), y_pred)

      # Get the total loss
      (scaled_loss, metric_loss,
       loss_metrics) = self.build_losses(y_pred['raw_output'], label)

      # Scale the loss for numerical stability
      if isinstance(optimizer, tf_keras.mixed_precision.LossScaleOptimizer):
        scaled_loss = optimizer.get_scaled_loss(scaled_loss)

    # Compute the gradient
    train_vars = model.trainable_variables
    gradients = tape.gradient(scaled_loss, train_vars)

    # Get unscaled loss if we are using the loss scale optimizer on fp16
    if isinstance(optimizer, tf_keras.mixed_precision.LossScaleOptimizer):
      gradients = optimizer.get_unscaled_gradients(gradients)

    # Apply gradients to the model
    optimizer.apply_gradients(zip(gradients, train_vars))
    logs = {self.loss: metric_loss}

    # Compute all metrics
    if metrics:
      for m in metrics:
        m.update_state(loss_metrics[m.name])
        logs.update({m.name: m.result()})
    return logs
//  [python]
def train_step(self, inputs, model, optimizer, metrics=None):
    """Train Step.

    Forward step and backwards propagate the model.

    Args:
      inputs: a dictionary of input tensors.
      model: the model, forward pass definition.
      optimizer: the optimizer for this training step.
      metrics: a nested structure of metrics objects.

    Returns:
      A dictionary of logs.
    """
    image, label = inputs

    with tf.GradientTape(persistent=False) as tape:
      # Compute a prediction
      y_pred = model(image, training=True)

      # Cast to float32 for gradietn computation
      y_pred = tf.nest.map_structure(lambda x: tf.cast(x, tf.float32), y_pred)

      # Get the total loss
      (scaled_loss, metric_loss,
       loss_metrics) = self.build_losses(y_pred['raw_output'], label)

      # Scale the loss for numerical stability
      if isinstance(optimizer, tf_keras.mixed_precision.LossScaleOptimizer):
        scaled_loss = optimizer.get_scaled_loss(scaled_loss)

    # Compute the gradient
    train_vars = model.trainable_variables
    gradients = tape.gradient(scaled_loss, train_vars)

    # Get unscaled loss if we are using the loss scale optimizer on fp16
    if isinstance(optimizer, tf_keras.mixed_precision.LossScaleOptimizer):
      gradients = optimizer.get_unscaled_gradients(gradients)

    # Apply gradients to the model
    optimizer.apply_gradients(zip(gradients, train_vars))
    logs = {self.loss: metric_loss}

    # Compute all metrics
    if metrics:
      for m in metrics:
        m.update_state(loss_metrics[m.name])
        logs.update({m.name: m.result()})
    return logs
//  [python]
def _gelu(x):
    """
    Gaussian Error Linear Unit. Original Implementation of the gelu activation function in Google Bert repo when
    initially created. For information: OpenAI GPT's gelu is slightly different (and gives slightly different results):
    0.5 * x * (1 + torch.tanh(math.sqrt(2 / math.pi) * (x + 0.044715 * torch.pow(x, 3)))) Also see
    https://arxiv.org/abs/1606.08415
    """
    x = tf.convert_to_tensor(x)
    cdf = 0.5 * (1.0 + tf.math.erf(x / tf.to_float(tf.sqrt(2.0))))

    return x * cdf
//  [python]
def _gelu(x):
    """
    Gaussian Error Linear Unit. Original Implementation of the gelu activation function in Google Bert repo when
    initially created. For information: OpenAI GPT's gelu is slightly different (and gives slightly different results):
    0.5 * x * (1 + torch.tanh(math.sqrt(2 / math.pi) * (x + 0.044715 * torch.pow(x, 3)))) Also see
    https://arxiv.org/abs/1606.08415
    """
    x = tf.convert_to_tensor(x)
    cdf = 0.5 * (1.0 + tf.math.erf(x / tf.cast(tf.sqrt(2.0), x.dtype)))

    return x * cdf
//  [python]
def validation_step(self, inputs, model, metrics=None):
    """Validatation step.

    Args:
      inputs: a dictionary of input tensors.
      model: the keras.Model.
      metrics: a nested structure of metrics objects.

    Returns:
      A dictionary of logs.
    """
    image, label = inputs

    # Step the model once
    y_pred = model(image, training=False)
    y_pred = tf.nest.map_structure(lambda x: tf.cast(x, tf.float32), y_pred)
    (_, metric_loss, loss_metrics) = self.build_losses(y_pred['raw_output'],
                                                       label)
    logs = {self.loss: metric_loss}

    # Reorganize and rescale the boxes
    info = label['groundtruths']['image_info']
    boxes = self._reorg_boxes(y_pred['bbox'], info, y_pred['num_detections'])

    # Build the input for the coc evaluation metric
    coco_model_outputs = {
        'detection_boxes': boxes,
        'detection_scores': y_pred['confidence'],
        'detection_classes': y_pred['classes'],
        'num_detections': y_pred['num_detections'],
        'source_id': label['groundtruths']['source_id'],
        'image_info': label['groundtruths']['image_info']
    }

    # Compute all metrics
    if metrics:
      logs.update(
          {self.coco_metric.name: (label['groundtruths'], coco_model_outputs)})
      for m in metrics:
        m.update_state(loss_metrics[m.name])
        logs.update({m.name: m.result()})
    return logs
//  [python]
def _gelu_new(x):
    """
    Gaussian Error Linear Unit. This is a smoother version of the GELU. Original paper: https://arxiv.org/abs/1606.0841

    Args:
        x: float Tensor to perform activation

    Returns:
        `x` with the GELU activation applied.
    """
    x = t.convert_to_tensor(x)
    pi = t.cast(math.pi, x.dtype)
    coeff = t.cast(0.044715, x.dtype)
    cdf = 0.5 * (1.0 + tf.tanh(tf.sqrt(2.0 / pi) * (x + coeff * tf.pow(x, 3))))

    return x * cdf
//  [python]
def _gelu_new(x):
    """
    Gaussian Error Linear Unit. This is a smoother version of the GELU. Original paper: https://arxiv.org/abs/1606.0841

    Args:
        x: float Tensor to perform activation

    Returns:
        `x` with the GELU activation applied.
    """
    x = tf.convert_to_tensor(x)
    pi = tf.cast(math.pi, x.dtype)
    coeff = tf.cast(0.044715, x.dtype)
    cdf = 0.5 * (1.0 + tf.tanh(tf.sqrt(2.0 / pi) * (x + coeff * tf.pow(x, 3))))

    return x * cdf
//  [python]
def get_losses_and_metrics(model, n_train):
  """Define the losses and metrics for the model."""

  def negative_log_likelihood(y, rv_y):
    del rv_y  # unused arg
    return -model.output.distribution.log_prob(tf.squeeze(y))

  def accuracy(y_true, y_sample):
    del y_sample  # unused arg
    return tf.equal(
        tf.argmax(input=model.output.distribution.logits, axis=1),
        tf.cast(tf.squeeze(y_true), tf.int64))

  def log_likelihood(y_true, y_sample):
    """Expected conditional log-likelihood."""
    del y_sample  # unused arg
    return model.output.distribution.log_prob(tf.squeeze(y_true))

  def kl(y_true, y_sample):
    """KL-divergence."""
    del y_true  # unused arg
    del y_sample  # unused arg
    sampling_cost = sum(
        [l.kl_cost_weight + l.kl_cost_bias for l in model.layers])
    return sum(model.losses) * n_train + sampling_cost

  def elbo(y_true, y_sample):
    return log_likelihood(y_true, y_sample) * n_train - kl(y_true, y_sample)

  return negative_log_likelihood, accuracy, log_likelihood, kl, elbo
//  [python]
def gelu_fast(x):
    x = t.convert_to_tensor(x)
    coeff1 = t.cast(0.044715, x.dtype)
    coeff2 = t.cast(0.7978845608, x.dtype)

    return 0.5 * x * (1.0 + t.tanh(x * coeff2 * (1.0 + coeff1 * x * x)))
//  [python]
def gelu_fast(x):
    x = tf.convert_to_tensor(x)
    coeff1 = tf.cast(0.044715, x.dtype)
    coeff2 = tf.cast(0.7978845608, x.dtype)

    return 0.5 * x * (1.0 + tf.tanh(x * coeff2 * (1.0 + coeff1 * x * x)))
//  [python]
def quick_gelu(x):
    x = tf.convert_to_tensor(x)
    coeff = tf.cast(1.702, x.dtype)
    return x * tf.math.sigmoid(coeff * x)
//  [python]
def _getSequenceLength(transposedInput):
        ''' return the sequence length of the batch on the premise that all non used values are filled with Nones / NaNs.
            transposed input must be of size (batch size, sequence length, number of features)

            IN      transposedInput     (tf.Tensor)     input sequence transposed to the above mentioned size
            OUT     length              (array)         length of the sequence for each state in the batch
        '''

        boolMaskNoneValues = tf.reduce_any(tf.is_nan(transposedInput), 2)
        usedPositions = tf.to_int32(tf.logical_not(boolMaskNoneValues))
        length = tf.reduce_sum(usedPositions, 1)

        return length
//  [python]
def _getSequenceLength(transposedInput):
        ''' return the sequence length of the batch on the premise that all non used values are filled with Nones / NaNs.
            transposed input must be of size (batch size, sequence length, number of features)

            IN      transposedInput     (tf.Tensor)     input sequence transposed to the above mentioned size
            OUT     length              (array)         length of the sequence for each state in the batch
        '''

        boolMaskNoneValues = tf.reduce_any(tf.is_nan(transposedInput), 2)
        usedPositions = tf.cast(tf.logical_not(boolMaskNoneValues), tf.int32)
        length = tf.reduce_sum(usedPositions, 1)

        return length
//  [python]
def _call_dense_pseudo_layer(self, inputs, params):
        """
        Calls an internal dense 'pseudo-layer' described in the build() documentation.
        """
        kernel, bias, activation, target_shape = params
        inputs = t.cast(inputs, self._compute_dtype)

        if K.is_sparse(inputs):
            outputs = sparse_ops.sparse_tensor_dense_matmul(inputs, kernel)
        else:
            outputs = gen_math_ops.mat_mul(inputs, kernel)

        if bias is not None:
            outputs = tf.nn.bias_add(outputs, bias)

        if activation is not None:
            outputs = activation(outputs)

        return tf.reshape(outputs, (-1, *target_shape))
//  [python]
def _call_dense_pseudo_layer(self, inputs, params):
        """
        Calls an internal dense 'pseudo-layer' described in the build() documentation.
        """
        kernel, bias, activation, target_shape = params
        inputs = tf.cast(inputs, self._compute_dtype)

        if K.is_sparse(inputs):
            outputs = sparse_ops.sparse_tensor_dense_matmul(inputs, kernel)
        else:
            outputs = gen_math_ops.mat_mul(inputs, kernel)

        if bias is not None:
            outputs = tf.nn.bias_add(outputs, bias)

        if activation is not None:
            outputs = activation(outputs)

        return tf.reshape(outputs, (-1, *target_shape))
//  [python]
def _mean_update(pre_mean, pre_count, x, pre_cap=None):

    # compute this batch stats
    this_sum = tf.reduce_sum(x, 0)
    this_bs = tf.cast(K.shape(x)[0], 'float32')  # this batch size

    # increase count and compute weights
    new_count = pre_count + this_bs
    alpha = this_bs / K.minimum(new_count, pre_cap)

    # compute new mean. Note that once we reach self.cap (e.g. 1000),
    # the 'previous mean' matters less
    new_mean = pre_mean * (1 - alpha) + (this_sum / this_bs) * alpha

    return (new_mean, new_count)
//  [python]
def draw_thresh(self, bounds, no_clip_tensor, prob, shape):
        assert 0 <= prob <= 1, f'{prob} is not a probability'
        assert tf.is_tensor(no_clip_tensor), 'not a tensor'

        def seed():
            return self.rand.integers(np.iinfo(int).max)

        # Threshold.
        if bounds is None or prob == 0:
            return no_clip_tensor

        if np.isscalar(bounds):
            clip_at = tf.constant(bounds, no_clip_tensor.dtype)

        else:
            clip_at = tf.random.uniform(shape, minval=bounds[0], maxval=bounds[1], seed=seed())
            if clip_at.dtype != no_clip_tensor.dtype:
                clip_at = t.cast(clip_at, no_clip_tensor.dtype)

        # Randomization.
        if prob < 1:
            rand_bit = tf.less(tf.random.uniform(shape, seed=seed()), prob)
            rand_bit = t.cast(rand_bit, no_clip_tensor.dtype)
            clip_at = rand_bit * clip_at + (1 - rand_bit) * no_clip_tensor

        return clip_at
//  [python]
def draw_thresh(self, bounds, no_clip_tensor, prob, shape):
        assert 0 <= prob <= 1, f'{prob} is not a probability'
        assert tf.is_tensor(no_clip_tensor), 'not a tensor'

        def seed():
            return self.rand.integers(np.iinfo(int).max)

        # Threshold.
        if bounds is None or prob == 0:
            return no_clip_tensor

        if np.isscalar(bounds):
            clip_at = tf.constant(bounds, no_clip_tensor.dtype)

        else:
            clip_at = tf.random.uniform(shape, minval=bounds[0], maxval=bounds[1], seed=seed())
            if clip_at.dtype != no_clip_tensor.dtype:
                clip_at = tf.cast(clip_at, no_clip_tensor.dtype)

        # Randomization.
        if prob < 1:
            rand_bit = tf.less(tf.random.uniform(shape, seed=seed()), prob)
            rand_bit = tf.cast(rand_bit, no_clip_tensor.dtype)
            clip_at = rand_bit * clip_at + (1 - rand_bit) * no_clip_tensor

        return clip_at
//  [python]
def _call_batch(self, inputs):
        """
        Performs dense operation on a single batch of input features,
        kernels weights, and optional bias weights.
        """
        x = tf.expand_dims(inputs[0], axis=0)  # add batch axis for input layer
        x = tf.cast(x, self._compute_dtype)
        kernel = inputs[1]

        if K.is_sparse(x):
            outputs = sparse_ops.sparse_tensor_dense_matmul(x, kernel)
        else:
            outputs = gen_math_ops.mat_mul(x, kernel)

        if self.use_bias:
            outputs = tf.nn.bias_add(outputs, inputs[2])

        outputs = outputs[0]  # remove added batch axis
        return outputs
//  [python]
def lenet5(n_examples, input_shape, num_classes):
  """Builds Bayesian LeNet5."""
  p_fn, q_fn = mean_field_fn(empirical_bayes=True)
  def normalized_kl_fn(q, p, _):
    return q.kl_divergence(p) / t.cast(n_examples, t.float32)

  inputs = tf.keras.layers.Input(shape=input_shape)
  conv1 = tfp.layers.Convolution2DFlipout(
      6,
      kernel_size=5,
      padding='SAME',
      activation=tf.nn.relu,
      kernel_prior_fn=p_fn,
      kernel_posterior_fn=q_fn,
      bias_prior_fn=p_fn,
      bias_posterior_fn=q_fn,
      kernel_divergence_fn=normalized_kl_fn,
      bias_divergence_fn=normalized_kl_fn)(inputs)
  pool1 = tf.keras.layers.MaxPooling2D(pool_size=[2, 2],
                                       strides=[2, 2],
                                       padding='SAME')(conv1)
  conv2 = tfp.layers.Convolution2DFlipout(
      16,
      kernel_size=5,
      padding='SAME',
      activation=tf.nn.relu,
      kernel_prior_fn=p_fn,
      kernel_posterior_fn=q_fn,
      bias_prior_fn=p_fn,
      bias_posterior_fn=q_fn,
      kernel_divergence_fn=normalized_kl_fn,
      bias_divergence_fn=normalized_kl_fn)(pool1)
  pool2 = tf.keras.layers.MaxPooling2D(pool_size=[2, 2],
                                       strides=[2, 2],
                                       padding='SAME')(conv2)
  conv3 = tfp.layers.Convolution2DFlipout(
      120,
      kernel_size=5,
      padding='SAME',
      activation=tf.nn.relu,
      kernel_prior_fn=p_fn,
      kernel_posterior_fn=q_fn,
      bias_prior_fn=p_fn,
      bias_posterior_fn=q_fn,
      kernel_divergence_fn=normalized_kl_fn,
      bias_divergence_fn=normalized_kl_fn)(pool2)
  flatten = tf.keras.layers.Flatten()(conv3)
  dense1 = tfp.layers.DenseLocalReparameterization(
      84,
      activation=tf.nn.relu,
      kernel_prior_fn=p_fn,
      kernel_posterior_fn=q_fn,
      bias_prior_fn=p_fn,
      bias_posterior_fn=q_fn,
      kernel_divergence_fn=normalized_kl_fn,
      bias_divergence_fn=normalized_kl_fn)(flatten)
  dense2 = tfp.layers.DenseLocalReparameterization(
      num_classes,
      kernel_prior_fn=p_fn,
      kernel_posterior_fn=q_fn,
      bias_prior_fn=p_fn,
      bias_posterior_fn=q_fn,
      kernel_divergence_fn=normalized_kl_fn,
      bias_divergence_fn=normalized_kl_fn)(dense1)
  outputs = tf.keras.layers.Lambda(lambda x: ed.Categorical(logits=x))(dense2)
  return tf.keras.models.Model(inputs=inputs, outputs=outputs)
//  [python]
def lenet5(n_examples, input_shape, num_classes):
  """Builds Bayesian LeNet5."""
  p_fn, q_fn = mean_field_fn(empirical_bayes=True)
  def normalized_kl_fn(q, p, _):
    return q.kl_divergence(p) / tf.cast(n_examples, tf.float32)

  inputs = tf.keras.layers.Input(shape=input_shape)
  conv1 = tfp.layers.Convolution2DFlipout(
      6,
      kernel_size=5,
      padding='SAME',
      activation=tf.nn.relu,
      kernel_prior_fn=p_fn,
      kernel_posterior_fn=q_fn,
      bias_prior_fn=p_fn,
      bias_posterior_fn=q_fn,
      kernel_divergence_fn=normalized_kl_fn,
      bias_divergence_fn=normalized_kl_fn)(inputs)
  pool1 = tf.keras.layers.MaxPooling2D(pool_size=[2, 2],
                                       strides=[2, 2],
                                       padding='SAME')(conv1)
  conv2 = tfp.layers.Convolution2DFlipout(
      16,
      kernel_size=5,
      padding='SAME',
      activation=tf.nn.relu,
      kernel_prior_fn=p_fn,
      kernel_posterior_fn=q_fn,
      bias_prior_fn=p_fn,
      bias_posterior_fn=q_fn,
      kernel_divergence_fn=normalized_kl_fn,
      bias_divergence_fn=normalized_kl_fn)(pool1)
  pool2 = tf.keras.layers.MaxPooling2D(pool_size=[2, 2],
                                       strides=[2, 2],
                                       padding='SAME')(conv2)
  conv3 = tfp.layers.Convolution2DFlipout(
      120,
      kernel_size=5,
      padding='SAME',
      activation=tf.nn.relu,
      kernel_prior_fn=p_fn,
      kernel_posterior_fn=q_fn,
      bias_prior_fn=p_fn,
      bias_posterior_fn=q_fn,
      kernel_divergence_fn=normalized_kl_fn,
      bias_divergence_fn=normalized_kl_fn)(pool2)
  flatten = tf.keras.layers.Flatten()(conv3)
  dense1 = tfp.layers.DenseLocalReparameterization(
      84,
      activation=tf.nn.relu,
      kernel_prior_fn=p_fn,
      kernel_posterior_fn=q_fn,
      bias_prior_fn=p_fn,
      bias_posterior_fn=q_fn,
      kernel_divergence_fn=normalized_kl_fn,
      bias_divergence_fn=normalized_kl_fn)(flatten)
  dense2 = tfp.layers.DenseLocalReparameterization(
      num_classes,
      kernel_prior_fn=p_fn,
      kernel_posterior_fn=q_fn,
      bias_prior_fn=p_fn,
      bias_posterior_fn=q_fn,
      kernel_divergence_fn=normalized_kl_fn,
      bias_divergence_fn=normalized_kl_fn)(dense1)
  outputs = tf.keras.layers.Lambda(lambda x: ed.Categorical(logits=x))(dense2)
  return tf.keras.models.Model(inputs=inputs, outputs=outputs)
//  [python]
def call(self, x, training=None):
        training = _get_training_value(training, self.trainable)

        # get batch shape:
        this_bs_int = K.shape(x)[0]

        # prep for broadcasting :(
        p = tf.concat((K.reshape(this_bs_int, (1,)), K.shape(self.cov)), 0)
        z = tf.ones(p)

        # If calling in inference mode, use moving stats
        if training is False:
            return K.minimum(1., self.count / self.cap) * (z * K.expand_dims(self.cov, 0))

        x_orig = x

        # update mean
        new_mean, new_count = _mean_update(self.mean, self.count, x, self.cap)

        # x reshape
        this_bs = t.cast(this_bs_int, 'float32')  # this batch size
        prev_count = self.count
        x = K.batch_flatten(x)  # B x N

        # new C update. Should be B x N x N
        x = K.expand_dims(x, -1)
        C_delta = K.batch_dot(x, K.permute_dimensions(x, [0, 2, 1]))

        # update cov
        prev_cap = K.minimum(prev_count, self.cap)
        C = self.cov * (prev_cap - 1) + K.sum(C_delta, 0)
        new_cov = C / (prev_cap + this_bs - 1)

        # updates
        self.count.assign(new_count)
        self.mean.assign(new_mean)
        self.cov.assign(new_cov)

        return K.minimum(1., new_count / self.cap) * (z * K.expand_dims(new_cov, 0))
//  [python]
def call(self, x, training=None):
        training = _get_training_value(training, self.trainable)

        # get batch shape:
        this_bs_int = K.shape(x)[0]

        # prep for broadcasting :(
        p = tf.concat((K.reshape(this_bs_int, (1,)), K.shape(self.cov)), 0)
        z = tf.ones(p)

        # If calling in inference mode, use moving stats
        if training is False:
            return K.minimum(1., self.count / self.cap) * (z * K.expand_dims(self.cov, 0))

        x_orig = x

        # update mean
        new_mean, new_count = _mean_update(self.mean, self.count, x, self.cap)

        # x reshape
        this_bs = tf.cast(this_bs_int, 'float32')  # this batch size
        prev_count = self.count
        x = K.batch_flatten(x)  # B x N

        # new C update. Should be B x N x N
        x = K.expand_dims(x, -1)
        C_delta = K.batch_dot(x, K.permute_dimensions(x, [0, 2, 1]))

        # update cov
        prev_cap = K.minimum(prev_count, self.cap)
        C = self.cov * (prev_cap - 1) + K.sum(C_delta, 0)
        new_cov = C / (prev_cap + this_bs - 1)

        # updates
        self.count.assign(new_count)
        self.mean.assign(new_mean)
        self.cov.assign(new_cov)

        return K.minimum(1., new_count / self.cap) * (z * K.expand_dims(new_cov, 0))
//  [python]
def __iou_tile_coordinate(x, tile_size):
        """Replicate a number across a bitmap of size tile_size"""

        xx = t.cast(t.round(x), dtype=t.int16)
        xx = t.expand_dims(xx, axis=-1)
        xx = t.tile(xx, [1, 1, tile_size])
        xx = t.expand_dims(xx, axis=2)
        xx = t.tile(xx, [1, 1, tile_size, 1])
        return xx
//  [python]
def __iou_tile_coordinate(x, tile_size):
        """Replicate a number across a bitmap of size tile_size"""

        xx = tf.cast(tf.round(x), dtype=tf.int16)
        xx = tf.expand_dims(xx, axis=-1)
        xx = tf.tile(xx, [1, 1, tile_size])
        xx = tf.expand_dims(xx, axis=2)
        xx = tf.tile(xx, [1, 1, tile_size, 1])
        return xx
//  [python]
def __iou_gen_linmap(batch, n, tile_size):
        """Creates two bitmaps filled with numbers increasing in X and Y direction.
        This trick makes it easier to draw filled rectangles usinf tf.less and tf.greater."""

        row = t.cast(t.linspace(0.0, (tile_size - 1) * 1.0, tile_size), dtype=t.int16)
        linmap = t.tile([row], [tile_size, 1])
        linmap = t.tile([linmap], [n, 1, 1])
        linmap = t.tile([linmap], [batch, 1, 1, 1])  # shape [batch, n, SIZE, SIZE]
        return linmap
//  [python]
def __iou_gen_linmap(batch, n, tile_size):
        """Creates two bitmaps filled with numbers increasing in X and Y direction.
        This trick makes it easier to draw filled rectangles usinf tf.less and tf.greater."""

        row = tf.cast(tf.linspace(0.0, (tile_size - 1) * 1.0, tile_size), dtype=tf.int16)
        linmap = tf.tile([row], [tile_size, 1])
        linmap = tf.tile([linmap], [n, 1, 1])
        linmap = tf.tile([linmap], [batch, 1, 1, 1])  # shape [batch, n, SIZE, SIZE]
        return linmap
//  [python]
def hungarian_matching(weights):
  """Computes the minimum linear sum assignment using the Hungarian algorithm.

  Args:
    weights: A float32 [batch_size, num_elems, num_elems] tensor, where each
      inner matrix represents weights to be use for matching.

  Returns:
    A bool [batch_size, num_elems, num_elems] tensor, where each element of the
    inner matrix represents whether the worker has been matched to the job.
    The returned matching will always be a perfect match.
  """
  batch_size, num_elems, _ = tf_utils.get_shape_list(weights, 3)

  weights = _prepare(weights)
  adj_matrix = tf.equal(weights, 0.)
  state, assignment = _maximum_bipartite_matching(adj_matrix)
  workers_cover, jobs_cover = _compute_cover(state, assignment)

  def _cover_incomplete(workers_cover, jobs_cover, *args):
    del args
    cover_sum = (
        tf.reduce_sum(t.int32(workers_cover)) +
        tf.reduce_sum(t.int32(jobs_cover)))
    return tf.less(cover_sum, batch_size * num_elems)

  def _update_weights_and_match(workers_cover, jobs_cover, weights, assignment):
    weights = _update_weights_using_cover(workers_cover, jobs_cover, weights)
    adj_matrix = tf.equal(weights, 0.)
    state, assignment = _maximum_bipartite_matching(adj_matrix, assignment)
    workers_cover, jobs_cover = _compute_cover(state, assignment)
    return workers_cover, jobs_cover, weights, assignment

  workers_cover, jobs_cover, weights, assignment = tf.while_loop(
      _cover_incomplete,
      _update_weights_and_match,
      (workers_cover, jobs_cover, weights, assignment),
      back_prop=False)
  return weights, assignment
//  [python]
def hungarian_matching(weights):
  """Computes the minimum linear sum assignment using the Hungarian algorithm.

  Args:
    weights: A float32 [batch_size, num_elems, num_elems] tensor, where each
      inner matrix represents weights to be use for matching.

  Returns:
    A bool [batch_size, num_elems, num_elems] tensor, where each element of the
    inner matrix represents whether the worker has been matched to the job.
    The returned matching will always be a perfect match.
  """
  batch_size, num_elems, _ = tf_utils.get_shape_list(weights, 3)

  weights = _prepare(weights)
  adj_matrix = tf.equal(weights, 0.)
  state, assignment = _maximum_bipartite_matching(adj_matrix)
  workers_cover, jobs_cover = _compute_cover(state, assignment)

  def _cover_incomplete(workers_cover, jobs_cover, *args):
    del args
    cover_sum = (
        tf.reduce_sum(tf.cast(workers_cover, tf.int32)) +
        tf.reduce_sum(tf.cast(jobs_cover, tf.int32)))
    return tf.less(cover_sum, batch_size * num_elems)

  def _update_weights_and_match(workers_cover, jobs_cover, weights, assignment):
    weights = _update_weights_using_cover(workers_cover, jobs_cover, weights)
    adj_matrix = tf.equal(weights, 0.)
    state, assignment = _maximum_bipartite_matching(adj_matrix, assignment)
    workers_cover, jobs_cover = _compute_cover(state, assignment)
    return workers_cover, jobs_cover, weights, assignment

  workers_cover, jobs_cover, weights, assignment = tf.while_loop(
      _cover_incomplete,
      _update_weights_and_match,
      (workers_cover, jobs_cover, weights, assignment),
      back_prop=False)
  return weights, assignment
//  [python]
def build_losses(self,
                   labels,
                   model_outputs,
                   metrics,
                   aux_losses=None) -> tf.Tensor:
    with tf.name_scope('MaskedLMTask/losses'):
      metrics = dict([(metric.name, metric) for metric in metrics])
      lm_prediction_losses = tf_keras.losses.sparse_categorical_crossentropy(
          labels['masked_lm_ids'],
          tf.cast(model_outputs['mlm_logits'], tf.float32),
          from_logits=True)
      lm_label_weights = labels['masked_lm_weights']
      lm_numerator_loss = tf.reduce_sum(lm_prediction_losses *
                                        lm_label_weights)
      lm_denominator_loss = tf.reduce_sum(lm_label_weights)
      mlm_loss = tf.math.divide_no_nan(lm_numerator_loss, lm_denominator_loss)
      metrics['lm_example_loss'].update_state(mlm_loss)
      if 'next_sentence_labels' in labels:
        sentence_labels = labels['next_sentence_labels']
        sentence_outputs = tf.cast(
            model_outputs['next_sentence'], dtype=tf.float32)
        sentence_loss = tf.reduce_mean(
            tf_keras.losses.sparse_categorical_crossentropy(
                sentence_labels, sentence_outputs, from_logits=True))
        metrics['next_sentence_loss'].update_state(sentence_loss)
        total_loss = mlm_loss + sentence_loss
      else:
        total_loss = mlm_loss

      if aux_losses:
        total_loss += tf.add_n(aux_losses)
      return total_loss
//  [python]
def __iou_gen_rectmap(cls, linmap, rects, tile_size):
        """Draws filled rectangles"""

        x1, y1, x2, y2 = tf.unstack(rects, axis=-1)  # shapes [batch, n]
        x1tile = cls.__iou_tile_coordinate(x1, tile_size)
        x2tile = cls.__iou_tile_coordinate(x2, tile_size)
        y1tile = cls.__iou_tile_coordinate(y1, tile_size)
        y2tile = cls.__iou_tile_coordinate(y2, tile_size)
        zeros = tf.zeros_like(linmap, dtype=tf.uint8)
        ones = tf.ones_like(linmap, dtype=tf.uint8)
        mapx = tf.where(tf.greater_equal(linmap, x1tile), ones, zeros)
        mapx = tf.where(tf.less(linmap, x2tile), mapx, zeros)
        mapy = tf.where(tf.greater_equal(linmap, y1tile), ones, zeros)
        mapy = tf.where(tf.less(linmap, y2tile), mapy, zeros)
        mapy = tf.matrix_transpose(mapy)
        map = tf.logical_and(t.cast(mapx, tf.bool), t.cast(mapy, tf.bool))
        return map
//  [python]
def __iou_gen_rectmap(cls, linmap, rects, tile_size):
        """Draws filled rectangles"""

        x1, y1, x2, y2 = tf.unstack(rects, axis=-1)  # shapes [batch, n]
        x1tile = cls.__iou_tile_coordinate(x1, tile_size)
        x2tile = cls.__iou_tile_coordinate(x2, tile_size)
        y1tile = cls.__iou_tile_coordinate(y1, tile_size)
        y2tile = cls.__iou_tile_coordinate(y2, tile_size)
        zeros = tf.zeros_like(linmap, dtype=tf.uint8)
        ones = tf.ones_like(linmap, dtype=tf.uint8)
        mapx = tf.where(tf.greater_equal(linmap, x1tile), ones, zeros)
        mapx = tf.where(tf.less(linmap, x2tile), mapx, zeros)
        mapy = tf.where(tf.greater_equal(linmap, y1tile), ones, zeros)
        mapy = tf.where(tf.less(linmap, y2tile), mapy, zeros)
        mapy = tf.matrix_transpose(mapy)
        map = tf.logical_and(tf.cast(mapx, tf.bool), tf.cast(mapy, tf.bool))
        return map
//  [python]
def batch_mean(ious):
        """Computes the average IOU across a batch of IOUs
        IOUs of value 1 mean that the network correctly detected nothing when there was
        nothing to detect. To compute the average IOU, 1 values are eliminated. The result
        is the average IOU across all instances where either something was detected or
        there was something to detect. In the rare case where the result would be 0/0,
        the return value is 1 which is not really correct but should be rare and offset
        a further average of batch_mean() results only a little.

        Args:
            ious: shape[batch]

        Returns:
            mean IOU
        """
        correct_non_detections = tf.equal(ious, 1.0)
        other_detections = tf.logical_not(correct_non_detections)
        n = tf.reduce_sum(tf.to_float(other_detections))
        m = tf.reduce_sum(tf.where(correct_non_detections, tf.zeros_like(ious), ious))
        safe_n = tf.where(tf.equal(n, 0.0), tf.ones_like(n), n)
        safe_m = tf.where(tf.equal(n, 0.0), tf.ones_like(m), m)
        return safe_m/safe_n
//  [python]
def batch_mean(ious):
        """Computes the average IOU across a batch of IOUs
        IOUs of value 1 mean that the network correctly detected nothing when there was
        nothing to detect. To compute the average IOU, 1 values are eliminated. The result
        is the average IOU across all instances where either something was detected or
        there was something to detect. In the rare case where the result would be 0/0,
        the return value is 1 which is not really correct but should be rare and offset
        a further average of batch_mean() results only a little.

        Args:
            ious: shape[batch]

        Returns:
            mean IOU
        """
        correct_non_detections = tf.equal(ious, 1.0)
        other_detections = tf.logical_not(correct_non_detections)
        n = tf.reduce_sum(tf.cast(other_detections, tf.float32))
        m = tf.reduce_sum(tf.where(correct_non_detections, tf.zeros_like(ious), ious))
        safe_n = tf.where(tf.equal(n, 0.0), tf.ones_like(n), n)
        safe_m = tf.where(tf.equal(n, 0.0), tf.ones_like(m), m)
        return safe_m/safe_n