jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def _minibatch_subsample_fn(self, inputs):
    """Randomly samples anchors for one image.

    Args:
      inputs: a list of 2 inputs. First one is a tensor of shape [num_anchors,
        num_classes] indicating targets assigned to each anchor. Second one
        is a tensor of shape [num_anchors] indicating the class weight of each
        anchor.

    Returns:
      batch_sampled_indicator: bool tensor of shape [num_anchors] indicating
        whether the anchor should be selected for loss computation.
    """
    cls_targets, cls_weights = inputs
    if self._add_background_class:
      # Set background_class bits to 0 so that the positives_indicator
      # computation would not consider background class.
      background_class = tf.zeros_like(tf.slice(cls_targets, [0, 0], [-1, 1]))
      regular_class = tf.slice(cls_targets, [0, 1], [-1, -1])
      cls_targets = tf.concat([background_class, regular_class], 1)
    positives_indicator = tf.reduce_sum(cls_targets, axis=1)
    return self._random_example_sampler.subsample(
        tf.cast(cls_weights, tf.bool),
        batch_size=None,
        labels=tf.cast(positives_indicator, tf.bool))
//  [python]
def calculate_threshold_confusion_matrix(labels_mask, preds, num_thresh):
  """Calculates confusion matrix based on thresholds.

  Given labels mask, predictions, and number of thresholds, returns count
  for cell in confusion matrix.

  Args:
    labels_mask: tf.bool vector tensor when label was normal or
      anomalous.
    preds: Predicted anomaly labels.
    num_thresh: Number of anomaly thresholds to try in parallel grid search.

  Returns:
    Count for cell in confusion matrix.
  """
  count = tf.reduce_sum(
      input_tensor=tf.to_int64(
          x=tf.map_fn(
              fn=lambda threshold: tf.logical_and(
                  x=labels_mask,
                  y=preds[threshold, :]),
              elems=tf.range(start=0, limit=num_thresh, dtype=tf.int64),
              dtype=tf.bool)),
      axis=1)

  return count
//  [python]
def calculate_threshold_confusion_matrix(labels_mask, preds, num_thresh):
  """Calculates confusion matrix based on thresholds.

  Given labels mask, predictions, and number of thresholds, returns count
  for cell in confusion matrix.

  Args:
    labels_mask: tf.bool vector tensor when label was normal or
      anomalous.
    preds: Predicted anomaly labels.
    num_thresh: Number of anomaly thresholds to try in parallel grid search.

  Returns:
    Count for cell in confusion matrix.
  """
  count = tf.reduce_sum(
      input_tensor=tf.cast(
          x=tf.map_fn(
              fn=lambda threshold: tf.logical_and(
                  x=labels_mask,
                  y=preds[threshold, :]),
              elems=tf.range(start=0, limit=num_thresh, dtype=tf.int64),
              dtype=tf.bool),
          dtype=tf.int64),
      axis=1)

  return count
//  [python]
def maskrcnn_loss(mask_logits, fg_labels, fg_target_masks):
    """
    Args:
        mask_logits: #fg x #category x14x14
        fg_labels: #fg, in 1~#class
        fg_target_masks: #fgx14x14, int
    """
    num_fg = tf.size(fg_labels)
    indices = tf.stack([tf.range(num_fg), tf.to_int32(fg_labels) - 1], axis=1)  # #fgx2
    mask_logits = tf.gather_nd(mask_logits, indices)  # #fgx14x14
    mask_probs = tf.sigmoid(mask_logits)

    # add some training visualizations to tensorboard
    with tf.name_scope('mask_viz'):
        viz = tf.concat([fg_target_masks, mask_probs], axis=1)
        viz = tf.expand_dims(viz, 3)
        viz = tf.to_int64(viz * 255, name='viz')
        tf.summary.image('mask_truth|pred', viz, max_outputs=10)

    loss = tf.nn.sigmoid_cross_entropy_with_logits(
        labels=fg_target_masks, logits=mask_logits)
    loss = tf.reduce_mean(loss, name='maskrcnn_loss')

    pred_label = mask_probs > 0.5
    truth_label = fg_target_masks > 0.5
    accuracy = tf.reduce_mean(
        tf.to_float(tf.equal(pred_label, truth_label)),
        name='accuracy')
    pos_accuracy = tf.logical_and(
        tf.equal(pred_label, truth_label),
        tf.equal(truth_label, True))
    pos_accuracy = tf.reduce_mean(tf.to_float(pos_accuracy), name='pos_accuracy')
    fg_pixel_ratio = tf.reduce_mean(tf.to_float(truth_label), name='fg_pixel_ratio')

    add_moving_summary(loss, accuracy, fg_pixel_ratio, pos_accuracy)
    return loss
//  [python]
def maskrcnn_loss(mask_logits, fg_labels, fg_target_masks):
    """
    Args:
        mask_logits: #fg x #category x14x14
        fg_labels: #fg, in 1~#class
        fg_target_masks: #fgx14x14, int
    """
    num_fg = tf.size(fg_labels)
    indices = tf.stack([tf.range(num_fg), tf.to_int32(fg_labels) - 1], axis=1)  # #fgx2
    mask_logits = tf.gather_nd(mask_logits, indices)  # #fgx14x14
    mask_probs = tf.sigmoid(mask_logits)

    # add some training visualizations to tensorboard
    with tf.name_scope('mask_viz'):
        viz = tf.concat([fg_target_masks, mask_probs], axis=1)
        viz = tf.expand_dims(viz, 3)
        viz = tf.cast(viz * 255, tf.uint8, name='viz')
        tf.summary.image('mask_truth|pred', viz, max_outputs=10)

    loss = tf.nn.sigmoid_cross_entropy_with_logits(
        labels=fg_target_masks, logits=mask_logits)
    loss = tf.reduce_mean(loss, name='maskrcnn_loss')

    pred_label = mask_probs > 0.5
    truth_label = fg_target_masks > 0.5
    accuracy = tf.reduce_mean(
        tf.to_float(tf.equal(pred_label, truth_label)),
        name='accuracy')
    pos_accuracy = tf.logical_and(
        tf.equal(pred_label, truth_label),
        tf.equal(truth_label, True))
    pos_accuracy = tf.reduce_mean(tf.to_float(pos_accuracy), name='pos_accuracy')
    fg_pixel_ratio = tf.reduce_mean(tf.to_float(truth_label), name='fg_pixel_ratio')

    add_moving_summary(loss, accuracy, fg_pixel_ratio, pos_accuracy)
    return loss
//  [python]
def dot_based_interact_native(input, bottom_mlp_output):
  # Dot Based Interact of the "input" tensor
  concat_features = t.cast(input, t.float32)
  interactions = tf.matmul(concat_features, concat_features, transpose_b=True)
  ones = tf.ones_like(interactions, dtype=concat_features.dtype)
  upper_tri_mask = tf.linalg.band_part(ones, 0, -1)
  feature_dim = tf.shape(interactions)[-1]
  lower_tri_mask = ones - upper_tri_mask
  activations = tf.boolean_mask(interactions, lower_tri_mask)
  out_dim = feature_dim * (feature_dim - 1) // 2
  activations = tf.reshape(activations, shape=[-1, out_dim])

  # Top Concatenation of the bottom_mlp_output with the interactions
  bottom_mlp_output = t.cast(tf.squeeze(bottom_mlp_output, axis=1), t.float32)
  top_concat = tf.concat([bottom_mlp_output, activations], axis=1)

  # Zero Padding for performance in upstream ops
  padding = tf.zeros([concat_features.shape[0], 1])
  zero_padded = tf.concat([top_concat, padding], axis=1)

  return zero_padded
//  [python]
def dot_based_interact_native(input, bottom_mlp_output):
  # Dot Based Interact of the "input" tensor
  concat_features = tf.cast(input, tf.float32)
  interactions = tf.matmul(concat_features, concat_features, transpose_b=True)
  ones = tf.ones_like(interactions, dtype=concat_features.dtype)
  upper_tri_mask = tf.linalg.band_part(ones, 0, -1)
  feature_dim = tf.shape(interactions)[-1]
  lower_tri_mask = ones - upper_tri_mask
  activations = tf.boolean_mask(interactions, lower_tri_mask)
  out_dim = feature_dim * (feature_dim - 1) // 2
  activations = tf.reshape(activations, shape=[-1, out_dim])

  # Top Concatenation of the bottom_mlp_output with the interactions
  bottom_mlp_output = tf.cast(tf.squeeze(bottom_mlp_output, axis=1), tf.float32)
  top_concat = tf.concat([bottom_mlp_output, activations], axis=1)

  # Zero Padding for performance in upstream ops
  padding = tf.zeros([concat_features.shape[0], 1])
  zero_padded = tf.concat([top_concat, padding], axis=1)

  return zero_padded
//  [python]
def calculate_composite_classification_metrics(tp, fn, fp, tn, f_score_beta):
  """Calculates compositive classification metrics from the confusion matrix.

  Given variables for the confusion matrix and the value of beta for f-beta
  score, returns accuracy, precision, recall, and f-beta score composite
  metrics.

  Args:
    tp: tf.int64 variable tracking number of true positives at
      each possible anomaly threshold.
    fn: tf.int64 variable tracking number of false negatives at
      each possible anomaly threshold.
    fp: tf.int64 variable tracking number of false positives at
      each possible anomaly threshold.
    tn: tf.int64 variable tracking number of true negatives at
      each possible anomaly threshold.
    f_score_beta: Value of beta for f-beta score.

  Returns:
    Accuracy, precision, recall, and f-beta score composite metric tensors.
  """
  # time_shape = (num_time_anom_thresh,)
  # feat_shape = (num_feat_anom_thresh,)
  acc = t.cast(x=tp + tn, dtype=t.float64) \
    / t.cast(x=tp + fn + fp + tn, dtype=t.float64)
  tp_float64 = t.cast(x=tp, dtype=t.float64)
  pre = tp_float64 / t.cast(x=tp + fp, dtype=t.float64)
  rec = tp_float64 / t.cast(x=tp + fn, dtype=t.float64)
  f_beta_numerator = (1.0 + f_score_beta ** 2) * (pre * rec)
  f_beta_score = f_beta_numerator / (f_score_beta ** 2 * pre + rec)

  return acc, pre, rec, f_beta_score
//  [python]
def calculate_composite_classification_metrics(tp, fn, fp, tn, f_score_beta):
  """Calculates compositive classification metrics from the confusion matrix.

  Given variables for the confusion matrix and the value of beta for f-beta
  score, returns accuracy, precision, recall, and f-beta score composite
  metrics.

  Args:
    tp: tf.int64 variable tracking number of true positives at
      each possible anomaly threshold.
    fn: tf.int64 variable tracking number of false negatives at
      each possible anomaly threshold.
    fp: tf.int64 variable tracking number of false positives at
      each possible anomaly threshold.
    tn: tf.int64 variable tracking number of true negatives at
      each possible anomaly threshold.
    f_score_beta: Value of beta for f-beta score.

  Returns:
    Accuracy, precision, recall, and f-beta score composite metric tensors.
  """
  # time_shape = (num_time_anom_thresh,)
  # feat_shape = (num_feat_anom_thresh,)
  acc = tf.cast(x=tp + tn, dtype=tf.float64) \
    / tf.cast(x=tp + fn + fp + tn, dtype=tf.float64)
  tp_float64 = tf.cast(x=tp, dtype=tf.float64)
  pre = tp_float64 / tf.cast(x=tp + fp, dtype=tf.float64)
  rec = tp_float64 / tf.cast(x=tp + fn, dtype=tf.float64)
  f_beta_numerator = (1.0 + f_score_beta ** 2) * (pre * rec)
  f_beta_score = f_beta_numerator / (f_score_beta ** 2 * pre + rec)

  return acc, pre, rec, f_beta_score
//  [python]
def secondclassification_losses(labels, label_logits):
    """
    Args:
        labels: n,
        label_logits: nxC
    """
    if config.USE_COCO_AND_MAPILLARY:
        # replace void label by 1 to avoid out of range. Later mask out the contribution of the void label to the loss
        labels_without_void = tf.where(tf.equal(labels, config.VOID_LABEL), tf.ones_like(labels), labels)
    else:
        labels_without_void = labels
    label_loss = tf.nn.sparse_softmax_cross_entropy_with_logits(
        labels=labels_without_void, logits=label_logits)
    if config.USE_COCO_AND_MAPILLARY:
        no_void_label_float_mask = tf.to_float(tf.not_equal(labels, config.VOID_LABEL))
        label_loss_sum = tf.reduce_sum(label_loss * no_void_label_float_mask)
        n = tf.maximum(tf.reduce_sum(no_void_label_float_mask), 1)
        label_loss = tf.divide(label_loss_sum, n, name='label_loss')
    else:
        label_loss = tf.reduce_mean(label_loss, name='label_loss')

    # to make the handling of void label easier, we don't report these stats when a void label is used for now
    if not config.USE_COCO_AND_MAPILLARY:
        fg_inds = tf.where(labels > 0)[:, 0]
        num_fg = tf.size(fg_inds)
        with tf.name_scope('label_metrics'), tf.device('/cpu:0'):
            prediction = tf.argmax(label_logits, axis=1, name='label_prediction')
            correct = tf.to_float(tf.equal(prediction, labels))  # boolean/integer gather is unavailable on GPU
            accuracy = tf.reduce_mean(correct, name='accuracy')
            fg_label_pred = tf.argmax(tf.gather(label_logits, fg_inds), axis=1)
            num_zero = tf.reduce_sum(tf.to_int32(tf.equal(fg_label_pred, 0)), name='num_zero')
            false_negative = tf.truediv(num_zero, num_fg, name='false_negative')
            fg_accuracy = tf.reduce_mean(
                tf.gather(correct, fg_inds), name='fg_accuracy')
        add_moving_summary(label_loss, accuracy, fg_accuracy, false_negative)

    return label_loss
//  [python]
def secondclassification_losses(labels, label_logits):
    """
    Args:
        labels: n,
        label_logits: nxC
    """
    if config.USE_COCO_AND_MAPILLARY:
        # replace void label by 1 to avoid out of range. Later mask out the contribution of the void label to the loss
        labels_without_void = tf.where(tf.equal(labels, config.VOID_LABEL), tf.ones_like(labels), labels)
    else:
        labels_without_void = labels
    label_loss = tf.nn.sparse_softmax_cross_entropy_with_logits(
        labels=labels_without_void, logits=label_logits)
    if config.USE_COCO_AND_MAPILLARY:
        no_void_label_float_mask = tf.cast(tf.not_equal(labels, config.VOID_LABEL), tf.float32)
        label_loss_sum = tf.reduce_sum(label_loss * no_void_label_float_mask)
        n = tf.maximum(tf.reduce_sum(no_void_label_float_mask), 1)
        label_loss = tf.divide(label_loss_sum, n, name='label_loss')
    else:
        label_loss = tf.reduce_mean(label_loss, name='label_loss')

    # to make the handling of void label easier, we don't report these stats when a void label is used for now
    if not config.USE_COCO_AND_MAPILLARY:
        fg_inds = tf.where(labels > 0)[:, 0]
        num_fg = tf.size(fg_inds)
        with tf.name_scope('label_metrics'), tf.device('/cpu:0'):
            prediction = tf.argmax(label_logits, axis=1, name='label_prediction')
            correct = tf.to_float(tf.equal(prediction, labels))  # boolean/integer gather is unavailable on GPU
            accuracy = tf.reduce_mean(correct, name='accuracy')
            fg_label_pred = tf.argmax(tf.gather(label_logits, fg_inds), axis=1)
            num_zero = tf.reduce_sum(tf.to_int32(tf.equal(fg_label_pred, 0)), name='num_zero')
            false_negative = tf.truediv(num_zero, num_fg, name='false_negative')
            fg_accuracy = tf.reduce_mean(
                tf.gather(correct, fg_inds), name='fg_accuracy')
        add_moving_summary(label_loss, accuracy, fg_accuracy, false_negative)

    return label_loss
//  [python]
def random_apply(func, gunc, p, x):
  """Randomly apply function func to x with probability p otherwise function gunc."""
  return tf.cond(
      tf.less(
          tf.random.uniform([], minval=0, maxval=1, dtype=tf.float32),
          t.cast(p, tf.float32)), lambda: func(x), lambda: gunc(x))
//  [python]
def random_apply(func, gunc, p, x):
  """Randomly apply function func to x with probability p otherwise function gunc."""
  return tf.cond(
      tf.less(
          tf.random.uniform([], minval=0, maxval=1, dtype=tf.float32),
          tf.cast(p, tf.float32)), lambda: func(x), lambda: gunc(x))
//  [python]
def timestampify(start: tf.Tensor,
                 end: tf.Tensor,
                 duration: tf.Tensor,
                 abs_time_token: bool,
                 num_bins: int,
                 vocabulary_size: int,
                 time_format: str,
                 t: int = 1000000):
  """Tokenizes timestamps.

  Args:
    start: Start times of the events.
    end: End times of the events.
    duration: Duration of the video.
    abs_time_token: Whether to use absolute (vs relative) time tokens.
    num_bins: Number of quantization bins for time tokens.
    vocabulary_size: Number of text tokens.
    time_format: st for start-end or cd for center-duration.
    t: FPS * 1000000

  Returns:
    Tensor of start+end time tokens for each event.
  """

  if time_format == 'cd':
    timestamp = tf.stack([(start + end) // 2, end - start], axis=1)
  else:
    timestamp = tf.stack([start, end], axis=1)
  timestamp = tf.minimum(timestamp, duration)

  if not abs_time_token:  # relative time token
    max_offset = tf.cast(num_bins - 1, tf.float64)
    rel_timestamp = tf.math.divide(timestamp, duration)
    timestamp_token = tf.math.add(tf.math.multiply(rel_timestamp, max_offset),
                                  vocabulary_size)
    timestamp_token = tf.to_int32(timestamp_token)

  else:  # absolute time token
    timestamp_token = tf.math.add(
        tf.to_int32(timestamp / t), vocabulary_size)

  return timestamp_token
//  [python]
def timestampify(start: tf.Tensor,
                 end: tf.Tensor,
                 duration: tf.Tensor,
                 abs_time_token: bool,
                 num_bins: int,
                 vocabulary_size: int,
                 time_format: str,
                 t: int = 1000000):
  """Tokenizes timestamps.

  Args:
    start: Start times of the events.
    end: End times of the events.
    duration: Duration of the video.
    abs_time_token: Whether to use absolute (vs relative) time tokens.
    num_bins: Number of quantization bins for time tokens.
    vocabulary_size: Number of text tokens.
    time_format: st for start-end or cd for center-duration.
    t: FPS * 1000000

  Returns:
    Tensor of start+end time tokens for each event.
  """

  if time_format == 'cd':
    timestamp = tf.stack([(start + end) // 2, end - start], axis=1)
  else:
    timestamp = tf.stack([start, end], axis=1)
  timestamp = tf.minimum(timestamp, duration)

  if not abs_time_token:  # relative time token
    max_offset = tf.cast(num_bins - 1, tf.float64)
    rel_timestamp = tf.math.divide(timestamp, duration)
    timestamp_token = tf.math.add(tf.math.multiply(rel_timestamp, max_offset),
                                  vocabulary_size)
    timestamp_token = tf.cast(timestamp_token, tf.int32)

  else:  # absolute time token
    timestamp_token = tf.math.add(
        tf.cast(timestamp / t, tf.int32), vocabulary_size)

  return timestamp_token
//  [python]
def noise_span_to_unique_sentinel(tokens, noise_mask, vocabulary_size, seeds):
  """T5-style preprocessing."""
  del seeds

  prev_token_is_noise = tf.pad(noise_mask[:-1], [[1, 0]])

  first_noise_tokens = tf.logical_and(
      noise_mask, tf.logical_not(prev_token_is_noise))
  subsequent_noise_tokens = tf.logical_and(noise_mask, prev_token_is_noise)

  sentinel = sentinel_id(vocabulary_size) + 1 - tf.cumsum(
      tf.to_int32(first_noise_tokens))

  tokens = tf.where(first_noise_tokens, sentinel, tokens)
  return tf.boolean_mask(tokens, tf.logical_not(subsequent_noise_tokens))
//  [python]
def noise_span_to_unique_sentinel(tokens, noise_mask, vocabulary_size, seeds):
  """T5-style preprocessing."""
  del seeds

  prev_token_is_noise = tf.pad(noise_mask[:-1], [[1, 0]])

  first_noise_tokens = tf.logical_and(
      noise_mask, tf.logical_not(prev_token_is_noise))
  subsequent_noise_tokens = tf.logical_and(noise_mask, prev_token_is_noise)

  sentinel = sentinel_id(vocabulary_size) + 1 - tf.cumsum(
      tf.cast(first_noise_tokens, tokens.dtype))

  tokens = tf.where(first_noise_tokens, sentinel, tokens)
  return tf.boolean_mask(tokens, tf.logical_not(subsequent_noise_tokens))
//  [python]
def tensor(data, dtype=None):
    if isinstance(data, tf.Tensor):
        if dtype is None or data.dtype == dtype:
            return data
        else:
            return tf.cast(data, dtype=dtype)
    else:
        if isinstance(data, numbers.Number):
            data = [data]
        return tf.convert_to_tensor(data, dtype=dtype)
//  [python]
def LeNet(path,
          conv_channels=None,
          linear_layers=None,
          batch_norm=True,
          batch_size=128,
          num_threads=4,
          min_queue_examples=1000,
          mode="train"):

    # Data.
    _maybe_download_cifar10(path)

    # Read images and labels from disk.
    if mode == "train":
        filenames = [os.path.join(path, CIFAR10_FOLDER, "data_batch_{}.bin".format(i)) for i in xrange(1, 6)]
    elif mode == "test":
        filenames = [os.path.join(path, CIFAR10_FOLDER, "test_batch.bin")]
    else:
        raise ValueError("Mode {} not recognised".format(mode))
    depth = 3
    height = 32
    width = 32
    label_bytes = 1
    image_bytes = depth * height * width
    record_bytes = label_bytes + image_bytes
    reader = tf.FixedLengthRecordReader(record_bytes=record_bytes)
    _, record = reader.read(tf.train.string_input_producer(filenames))
    record_bytes = tf.decode_raw(record, tf.uint8)

    label = tf.to_int32(tf.slice(record_bytes, [0], [label_bytes]))
    raw_image = tf.slice(record_bytes, [label_bytes], [image_bytes])
    image = tf.to_float(tf.reshape(raw_image, [depth, height, width]))
    # height x width x depth.
    image = tf.transpose(image, [1, 2, 0])
    image = tf.div(image, 255)

    queue = tf.RandomShuffleQueue(capacity=min_queue_examples + 3 * batch_size,
                                  min_after_dequeue=min_queue_examples,
                                  dtypes=[tf.float32, tf.int32],
                                  shapes=[image.get_shape(), label.get_shape()])
    enqueue_ops = [queue.enqueue([image, label]) for _ in xrange(num_threads)]
    tf.train.add_queue_runner(tf.train.QueueRunner(queue, enqueue_ops))

    # Network.
    def _conv_activation(x):  # pylint: disable=invalid-name
        return tf.nn.max_pool(tf.sigmoid(x),
                              ksize=[1, 2, 2, 1],
                              strides=[1, 2, 2, 1],
                              padding="VALID")

    conv = snt.nets.ConvNet2D(output_channels=conv_channels,
                              kernel_shapes=[5],
                              strides=[1],
                              paddings=[snt.VALID],
                              activation=_conv_activation,
                              activate_final=True,
                              initializers=_nn_initializers,
                              use_batch_norm=batch_norm)

    if batch_norm:
        linear_activation = lambda x: tf.sigmoid(snt.BatchNorm()(x, is_training=True))
    else:
        linear_activation = tf.sigmoid

    mlp = snt.nets.MLP(list(linear_layers) + [10],
                       activation=linear_activation,
                       initializers=_nn_initializers)
    network = snt.Sequential([conv, snt.BatchFlatten(), mlp])

    def build():
        image_batch, label_batch = queue.dequeue_many(batch_size)
        label_batch = tf.reshape(label_batch, [batch_size])

        output = network(image_batch)
        return _xent_loss(output, label_batch)

    return build
//  [python]
def LeNet(path,
          conv_channels=None,
          linear_layers=None,
          batch_norm=True,
          batch_size=128,
          num_threads=4,
          min_queue_examples=1000,
          mode="train"):

    # Data.
    _maybe_download_cifar10(path)

    # Read images and labels from disk.
    if mode == "train":
        filenames = [os.path.join(path, CIFAR10_FOLDER, "data_batch_{}.bin".format(i)) for i in xrange(1, 6)]
    elif mode == "test":
        filenames = [os.path.join(path, CIFAR10_FOLDER, "test_batch.bin")]
    else:
        raise ValueError("Mode {} not recognised".format(mode))
    depth = 3
    height = 32
    width = 32
    label_bytes = 1
    image_bytes = depth * height * width
    record_bytes = label_bytes + image_bytes
    reader = tf.FixedLengthRecordReader(record_bytes=record_bytes)
    _, record = reader.read(tf.train.string_input_producer(filenames))
    record_bytes = tf.decode_raw(record, tf.uint8)

    label = tf.cast(tf.slice(record_bytes, [0], [label_bytes]), tf.int32)
    raw_image = tf.slice(record_bytes, [label_bytes], [image_bytes])
    image = tf.cast(tf.reshape(raw_image, [depth, height, width]), tf.float32)
    # height x width x depth.
    image = tf.transpose(image, [1, 2, 0])
    image = tf.div(image, 255)

    queue = tf.RandomShuffleQueue(capacity=min_queue_examples + 3 * batch_size,
                                  min_after_dequeue=min_queue_examples,
                                  dtypes=[tf.float32, tf.int32],
                                  shapes=[image.get_shape(), label.get_shape()])
    enqueue_ops = [queue.enqueue([image, label]) for _ in xrange(num_threads)]
    tf.train.add_queue_runner(tf.train.QueueRunner(queue, enqueue_ops))

    # Network.
    def _conv_activation(x):  # pylint: disable=invalid-name
        return tf.nn.max_pool(tf.sigmoid(x),
                              ksize=[1, 2, 2, 1],
                              strides=[1, 2, 2, 1],
                              padding="VALID")

    conv = snt.nets.ConvNet2D(output_channels=conv_channels,
                              kernel_shapes=[5],
                              strides=[1],
                              paddings=[snt.VALID],
                              activation=_conv_activation,
                              activate_final=True,
                              initializers=_nn_initializers,
                              use_batch_norm=batch_norm)

    if batch_norm:
        linear_activation = lambda x: tf.sigmoid(snt.BatchNorm()(x, is_training=True))
    else:
        linear_activation = tf.sigmoid

    mlp = snt.nets.MLP(list(linear_layers) + [10],
                       activation=linear_activation,
                       initializers=_nn_initializers)
    network = snt.Sequential([conv, snt.BatchFlatten(), mlp])

    def build():
        image_batch, label_batch = queue.dequeue_many(batch_size)
        label_batch = tf.reshape(label_batch, [batch_size])

        output = network(image_batch)
        return _xent_loss(output, label_batch)

    return build
//  [python]
def astype(input, ty):
    with tf.device(input.device):
        return tf.cast(input, dtype=ty)
//  [python]
def vgg16_cifar10(path,  # pylint: disable=invalid-name
            batch_norm=False,
            batch_size=128,
            num_threads=4,
            min_queue_examples=1000,
            mode="train"):
    """Cifar10 classification with a convolutional network."""
    
    # Data.
    _maybe_download_cifar10(path)
    # pdb.set_trace()
    # Read images and labels from disk.
    if mode == "train":
        filenames = [os.path.join(path,
                                  CIFAR10_FOLDER,
                                  "data_batch_{}.bin".format(i))
                     for i in xrange(1, 6)]
        is_training = True
    elif mode == "test":
        filenames = [os.path.join(path, CIFAR10_FOLDER, "test_batch.bin")]
        is_training = False
    else:
        raise ValueError("Mode {} not recognised".format(mode))
    
    depth = 3
    height = 32
    width = 32
    label_bytes = 1
    image_bytes = depth * height * width
    record_bytes = label_bytes + image_bytes
    reader = tf.FixedLengthRecordReader(record_bytes=record_bytes)
    _, record = reader.read(tf.train.string_input_producer(filenames))
    record_bytes = tf.decode_raw(record, tf.uint8)
    
    label = tf.cast(tf.slice(record_bytes, [0], [label_bytes]), tf.int32)
    raw_image = tf.slice(record_bytes, [label_bytes], [image_bytes])
    image = tf.cast(tf.reshape(raw_image, [depth, height, width]), tf.float32)
    # height x width x depth.
    image = tf.transpose(image, [1, 2, 0])
    image = tf.math.divide(image, 255)

    queue = tf.RandomShuffleQueue(capacity=min_queue_examples + 3 * batch_size,
                                  min_after_dequeue=min_queue_examples,
                                  dtypes=[tf.float32, tf.int32],
                                  shapes=[image.get_shape(), label.get_shape()])
    enqueue_ops = [queue.enqueue([image, label]) for _ in xrange(num_threads)]
    tf.train.add_queue_runner(tf.train.QueueRunner(queue, enqueue_ops))

    vgg = VGG16(0.5, 10)
    def build():
        image_batch, label_batch = queue.dequeue_many(batch_size)
        label_batch = tf.reshape(label_batch, [batch_size])
        # pdb.set_trace()
        output = vgg._build_model(image_batch)
        # print(output.shape)
        return _xent_loss(output, label_batch)
    
    return build
//  [python]
def sparse_matrix(data, index, shape, force_format=False):
    fmt = index[0]
    if fmt != "coo":
        raise TypeError(
            "Tensorflow backend only supports COO format. But got %s." % fmt
        )
    # tf.SparseTensor only supports int64 indexing,
    # therefore manually casting to int64 when input in int32
    spmat = tf.SparseTensor(
        indices=t.cast(tf.transpose(index[1], (1, 0)), tf.int64),
        values=data,
        dense_shape=shape,
    )
    return spmat, None
//  [python]
def sparse_matrix(data, index, shape, force_format=False):
    fmt = index[0]
    if fmt != "coo":
        raise TypeError(
            "Tensorflow backend only supports COO format. But got %s." % fmt
        )
    # tf.SparseTensor only supports int64 indexing,
    # therefore manually casting to int64 when input in int32
    spmat = tf.SparseTensor(
        indices=tf.cast(tf.transpose(index[1], (1, 0)), tf.int64),
        values=data,
        dense_shape=shape,
    )
    return spmat, None
//  [python]
def sum(input, dim, keepdims=False):
    if input.dtype == tf.bool:
        input = t.cast(input, tf.int32)
    return tf.reduce_sum(input, axis=dim, keepdims=keepdims)
//  [python]
def sum(input, dim, keepdims=False):
    if input.dtype == tf.bool:
        input = tf.cast(input, tf.int32)
    return tf.reduce_sum(input, axis=dim, keepdims=keepdims)
//  [python]
def reduce_sum(input):
    if input.dtype == tf.bool:
        input = t.cast(input, tf.int32)
    return tf.reduce_sum(input)
//  [python]
def reduce_sum(input):
    if input.dtype == tf.bool:
        input = tf.cast(input, tf.int32)
    return tf.reduce_sum(input)
//  [python]
def cumsum(input, dim):
    if input.dtype == tf.bool:
        input = t.cast(input, tf.int32)
    return tf.cumsum(input, axis=dim)
//  [python]
def cumsum(input, dim):
    if input.dtype == tf.bool:
        input = tf.cast(input, tf.int32)
    return tf.cumsum(input, axis=dim)
//  [python]
def argsort(input, dim, descending):
    if descending:
        return t.cast(
            tf.argsort(input, axis=dim, direction="DESCENDING"), dtype=tf.int64
        )
    else:
        return t.cast(
            tf.argsort(input, axis=dim, direction="ASCENDING"), dtype=tf.int64
        )
//  [python]
def argsort(input, dim, descending):
    if descending:
        return tf.cast(
            tf.argsort(input, axis=dim, direction="DESCENDING"), dtype=tf.int64
        )
    else:
        return tf.cast(
            tf.argsort(input, axis=dim, direction="ASCENDING"), dtype=tf.int64
        )
//  [python]
def full_1d(length, fill_value, dtype, ctx):
    with tf.device(ctx):
        t = tf.fill([length], value=fill_value)
        t = t.astype(dtype=dtype)
    return t
//  [python]
def full_1d(length, fill_value, dtype, ctx):
    with tf.device(ctx):
        t = tf.fill([length], value=fill_value)
        t = tf.cast(t, dtype=dtype)
    return t
//  [python]
def nonzero_1d(input):
    nonzero_bool = t.cast(input, t.bool)
    return tf.reshape(tf.where(nonzero_bool), (-1,))
//  [python]
def nonzero_1d(input):
    nonzero_bool = tf.cast(input, tf.bool)
    return tf.reshape(tf.where(nonzero_bool), (-1,))
//  [python]
def sort_1d(input):
    return tf.sort(input), tf.cast(tf.argsort(input), dtype=tf.int64)
//  [python]
def zerocopy_to_dgl_ndarray(data):
    if device_type(data.device) == "gpu" and data.dtype in (tf.int32, tf.int64):
        # NOTE: TF doesn't keep signed tensors on GPU due to legacy issues with
        #   shape inference. Convert it to unsigned and cast it back afterwards.
        if data.dtype == tf.int32:
            data = t.cast(data, tf.uint32)
        elif data.dtype == tf.int64:
            data = t.cast(data, tf.uint64)
        return nd.cast_to_signed(nd.from_dlpack(zerocopy_to_dlpack(data)))
    else:
        return nd.from_dlpack(zerocopy_to_dlpack(data))
//  [python]
def zerocopy_to_dgl_ndarray(data):
    if device_type(data.device) == "gpu" and data.dtype in (tf.int32, tf.int64):
        # NOTE: TF doesn't keep signed tensors on GPU due to legacy issues with
        #   shape inference. Convert it to unsigned and cast it back afterwards.
        if data.dtype == tf.int32:
            data = tf.cast(data, tf.uint32)
        elif data.dtype == tf.int64:
            data = tf.cast(data, tf.uint64)
        return nd.cast_to_signed(nd.from_dlpack(zerocopy_to_dlpack(data)))
    else:
        return nd.from_dlpack(zerocopy_to_dlpack(data))
//  [python]
def _get_mask_for(
        self,
        tf_batch_data: Dict[Text, Dict[Text, List[tf.Tensor]]],
        key: Text,
        sub_key: Text,
    ) -> Optional[tf.Tensor]:
        if key not in tf_batch_data or sub_key not in tf_batch_data[key]:
            return None

        sequence_lengths = tf.to_int32(tf_batch_data[key][sub_key][0])
        return rasa_layers.compute_mask(sequence_lengths)
//  [python]
def _get_mask_for(
        self,
        tf_batch_data: Dict[Text, Dict[Text, List[tf.Tensor]]],
        key: Text,
        sub_key: Text,
    ) -> Optional[tf.Tensor]:
        if key not in tf_batch_data or sub_key not in tf_batch_data[key]:
            return None

        sequence_lengths = tf.cast(tf_batch_data[key][sub_key][0], dtype=tf.int32)
        return rasa_layers.compute_mask(sequence_lengths)
//  [python]
def _get_sequence_feature_lengths(
        self, tf_batch_data: Dict[Text, Dict[Text, List[tf.Tensor]]], key: Text
    ) -> tf.Tensor:
        """Fetches the sequence lengths of real tokens per input example.

        The number of real tokens for an example is the same as the length of the
        sequence of the sequence-level (token-level) features for that input example.
        """
        if key in tf_batch_data and SEQUENCE_LENGTH in tf_batch_data[key]:
            return tf.to_int32(tf_batch_data[key][SEQUENCE_LENGTH][0])

        batch_dim = self._get_batch_dim(tf_batch_data[key])
        return tf.zeros([batch_dim], dtype=tf.int32)
//  [python]
def _get_sequence_feature_lengths(
        self, tf_batch_data: Dict[Text, Dict[Text, List[tf.Tensor]]], key: Text
    ) -> tf.Tensor:
        """Fetches the sequence lengths of real tokens per input example.

        The number of real tokens for an example is the same as the length of the
        sequence of the sequence-level (token-level) features for that input example.
        """
        if key in tf_batch_data and SEQUENCE_LENGTH in tf_batch_data[key]:
            return tf.cast(tf_batch_data[key][SEQUENCE_LENGTH][0], dtype=tf.int32)

        batch_dim = self._get_batch_dim(tf_batch_data[key])
        return tf.zeros([batch_dim], dtype=tf.int32)
//  [python]
def training_inputs(self, input_ids, input_mask):
    input_values = self.embedding(input_ids)
    if self.embedding_size != self.model_dimension:
      input_values = self.input_bottleneck(input_values, input_mask)
    pos_indices = tf.cumsum(input_mask, axis=1, exclusive=True)
    pos_indices = tf.to_int32(pos_indices)
    pos_values = self.positional_embedding(pos_indices)
    inputs = self.qact(self.ln(input_values + pos_values))
    return inputs
//  [python]
def training_inputs(self, input_ids, input_mask):
    input_values = self.embedding(input_ids)
    if self.embedding_size != self.model_dimension:
      input_values = self.input_bottleneck(input_values, input_mask)
    pos_indices = tf.cumsum(input_mask, axis=1, exclusive=True)
    pos_indices = tf.cast(pos_indices, dtype=tf.int32)
    pos_values = self.positional_embedding(pos_indices)
    inputs = self.qact(self.ln(input_values + pos_values))
    return inputs
//  [python]
def call(self, positive_matches: tf.Tensor, negative_matches: tf.Tensor,
           ignored_matches: tf.Tensor):
    """Samples and selects positive and negative instances.

    Args:
      positive_matches: A `bool` tensor of shape of [batch, N] where N is the
        number of instances. For each element, `True` means the instance
        corresponds to a positive example.
      negative_matches: A `bool` tensor of shape of [batch, N] where N is the
        number of instances. For each element, `True` means the instance
        corresponds to a negative example.
      ignored_matches: A `bool` tensor of shape of [batch, N] where N is the
        number of instances. For each element, `True` means the instance should
        be ignored.

    Returns:
      A `tf.tensor` of shape of [batch_size, K], storing the indices of the
        sampled examples, where K is `num_samples`.
    """
    sample_candidates = tf.logical_and(
        tf.logical_or(positive_matches, negative_matches),
        tf.logical_not(ignored_matches))

    sampler = sampling_ops.BalancedPositiveNegativeSampler(
        positive_fraction=self._config_dict['foreground_fraction'],
        is_static=True)

    batch_size = sample_candidates.shape[0]
    sampled_indicators = []
    for i in range(batch_size):
      sampled_indicator = sampler.subsample(
          sample_candidates[i],
          self._config_dict['num_samples'],
          positive_matches[i])
      sampled_indicators.append(sampled_indicator)
    sampled_indicators = tf.stack(sampled_indicators)
    _, selected_indices = tf.nn.top_k(
        tf.to_int32(sampled_indicators),
        k=self._config_dict['num_samples'],
        sorted=True)

    return selected_indices
//  [python]
def call(self, positive_matches: tf.Tensor, negative_matches: tf.Tensor,
           ignored_matches: tf.Tensor):
    """Samples and selects positive and negative instances.

    Args:
      positive_matches: A `bool` tensor of shape of [batch, N] where N is the
        number of instances. For each element, `True` means the instance
        corresponds to a positive example.
      negative_matches: A `bool` tensor of shape of [batch, N] where N is the
        number of instances. For each element, `True` means the instance
        corresponds to a negative example.
      ignored_matches: A `bool` tensor of shape of [batch, N] where N is the
        number of instances. For each element, `True` means the instance should
        be ignored.

    Returns:
      A `tf.tensor` of shape of [batch_size, K], storing the indices of the
        sampled examples, where K is `num_samples`.
    """
    sample_candidates = tf.logical_and(
        tf.logical_or(positive_matches, negative_matches),
        tf.logical_not(ignored_matches))

    sampler = sampling_ops.BalancedPositiveNegativeSampler(
        positive_fraction=self._config_dict['foreground_fraction'],
        is_static=True)

    batch_size = sample_candidates.shape[0]
    sampled_indicators = []
    for i in range(batch_size):
      sampled_indicator = sampler.subsample(
          sample_candidates[i],
          self._config_dict['num_samples'],
          positive_matches[i])
      sampled_indicators.append(sampled_indicator)
    sampled_indicators = tf.stack(sampled_indicators)
    _, selected_indices = tf.nn.top_k(
        tf.cast(sampled_indicators, dtype=tf.int32),
        k=self._config_dict['num_samples'],
        sorted=True)

    return selected_indices
//  [python]
def _calculate_entity_loss(
        self,
        inputs: tf.Tensor,
        tag_ids: tf.Tensor,
        mask: tf.Tensor,
        sequence_lengths: tf.Tensor,
        tag_name: Text,
        entity_tags: Optional[tf.Tensor] = None,
    ) -> Tuple[tf.Tensor, tf.Tensor, tf.Tensor]:

        tag_ids = tf.to_int32(tag_ids[:, :, 0])

        if entity_tags is not None:
            _tags = self._tf_layers[f"embed.{tag_name}.tags"](entity_tags)
            inputs = tf.concat([inputs, _tags], axis=-1)

        logits = self._tf_layers[f"embed.{tag_name}.logits"](inputs)

        # should call first to build weights
        pred_ids, _ = self._tf_layers[f"crf.{tag_name}"](logits, sequence_lengths)
        loss = self._tf_layers[f"crf.{tag_name}"].loss(
            logits, tag_ids, sequence_lengths
        )
        f1 = self._tf_layers[f"crf.{tag_name}"].f1_score(tag_ids, pred_ids, mask)

        return loss, f1, logits
//  [python]
def _calculate_entity_loss(
        self,
        inputs: tf.Tensor,
        tag_ids: tf.Tensor,
        mask: tf.Tensor,
        sequence_lengths: tf.Tensor,
        tag_name: Text,
        entity_tags: Optional[tf.Tensor] = None,
    ) -> Tuple[tf.Tensor, tf.Tensor, tf.Tensor]:

        tag_ids = tf.cast(tag_ids[:, :, 0], tf.int32)

        if entity_tags is not None:
            _tags = self._tf_layers[f"embed.{tag_name}.tags"](entity_tags)
            inputs = tf.concat([inputs, _tags], axis=-1)

        logits = self._tf_layers[f"embed.{tag_name}.logits"](inputs)

        # should call first to build weights
        pred_ids, _ = self._tf_layers[f"crf.{tag_name}"](logits, sequence_lengths)
        loss = self._tf_layers[f"crf.{tag_name}"].loss(
            logits, tag_ids, sequence_lengths
        )
        f1 = self._tf_layers[f"crf.{tag_name}"].f1_score(tag_ids, pred_ids, mask)

        return loss, f1, logits
//  [python]
def scale(boxlist, y_scale, x_scale, scope=None):
  """scale box coordinates in x and y dimensions.

  Args:
    boxlist: BoxList holding N boxes
    y_scale: (float) scalar tensor
    x_scale: (float) scalar tensor
    scope: name scope.

  Returns:
    boxlist: BoxList holding N boxes
  """
  with tf.name_scope(scope or 'Scale'):
    y_scale = tf.cast(y_scale, tf.float32)
    x_scale = tf.cast(x_scale, tf.float32)
    y_min, x_min, y_max, x_max = tf.split(
        value=boxlist.get(), num_or_size_splits=4, axis=1)
    y_min = y_scale * y_min
    y_max = y_scale * y_max
    x_min = x_scale * x_min
    x_max = x_scale * x_max
    scaled_boxlist = box_list.BoxList(
        tf.concat([y_min, x_min, y_max, x_max], 1))
    return _copy_extra_fields(scaled_boxlist, boxlist)
//  [python]
def visualize_boxes_in_image(image, boxlist, normalized=False, scope=None):
  """Overlay bounding box list on image.

  Currently this visualization plots a 1 pixel thick red bounding box on top
  of the image.  Note that tf.image.draw_bounding_boxes essentially is
  1 indexed.

  Args:
    image: an image tensor with shape [height, width, 3]
    boxlist: a BoxList
    normalized: (boolean) specify whether corners are to be interpreted as
      absolute coordinates in image space or normalized with respect to the
      image size.
    scope: name scope.

  Returns:
    image_and_boxes: an image tensor with shape [height, width, 3]
  """
  with tf.name_scope(scope or 'VisualizeBoxesInImage'):
    if not normalized:
      height, width, _ = tf.unstack(tf.shape(image))
      boxlist = scale(boxlist, 1.0 / t.cast(height, tf.float32),
                      1.0 / t.cast(width, tf.float32))
    corners = tf.expand_dims(boxlist.get(), 0)
    image = tf.expand_dims(image, 0)
    return tf.squeeze(tf.image.draw_bounding_boxes(image, corners), [0])
//  [python]
def visualize_boxes_in_image(image, boxlist, normalized=False, scope=None):
  """Overlay bounding box list on image.

  Currently this visualization plots a 1 pixel thick red bounding box on top
  of the image.  Note that tf.image.draw_bounding_boxes essentially is
  1 indexed.

  Args:
    image: an image tensor with shape [height, width, 3]
    boxlist: a BoxList
    normalized: (boolean) specify whether corners are to be interpreted as
      absolute coordinates in image space or normalized with respect to the
      image size.
    scope: name scope.

  Returns:
    image_and_boxes: an image tensor with shape [height, width, 3]
  """
  with tf.name_scope(scope or 'VisualizeBoxesInImage'):
    if not normalized:
      height, width, _ = tf.unstack(tf.shape(image))
      boxlist = scale(boxlist, 1.0 / tf.cast(height, tf.float32),
                      1.0 / tf.cast(width, tf.float32))
    corners = tf.expand_dims(boxlist.get(), 0)
    image = tf.expand_dims(image, 0)
    return tf.squeeze(tf.image.draw_bounding_boxes(image, corners), [0])
//  [python]
def _encode_tf(self, s: tf.Tensor) -> tf.Tensor:
    """Encode a string tf.Tensor.

    Args:
      s: a tf.Scalar or a 1d tf.Tensor in [batch_size] with dtype tf.string

    Returns:
      a 1d tf.Tensor or
      a tf.RaggedTensor in [batch_size, None] with dtype tf.int32
    """

    def _encode_py_func(text) -> Sequence[int]:
      text = tf.compat.as_text(text.numpy().decode("UTF-8"))
      result = self._tokenizer.encode(text=text)
      result = tf.expand_dims(result, 0)
      return tf.to_int32(result)

    def _encode_func():
      result = tf.py_function(func=_encode_py_func, inp=[s], Tout=tf.int32)
      result = tf.squeeze(result, 0)
      result.set_shape([None])
      return result

    def _batch_encode_py_func(strs) -> Sequence[Sequence[int]]:
      texts = []
      for s in strs.numpy():
        texts.append(tf.compat.as_text(s.decode("UTF-8")))
      results = self._tokenizer.batch_encode_plus(texts, padding=False)
      input_ids = results.input_ids
      return tf.ragged.constant(input_ids, dtype=tf.int32)

    def _batch_encode_func():
      return tf.py_function(
          func=_batch_encode_py_func,
          inp=[s],
          Tout=tf.RaggedTensorSpec(shape=[s.shape[0], None], dtype=tf.int32),
      )

    # pylint: disable=g-explicit-length-test
    if len(s.shape) > 0:
      return _batch_encode_func()
    return _encode_func()
//  [python]
def _encode_tf(self, s: tf.Tensor) -> tf.Tensor:
    """Encode a string tf.Tensor.

    Args:
      s: a tf.Scalar or a 1d tf.Tensor in [batch_size] with dtype tf.string

    Returns:
      a 1d tf.Tensor or
      a tf.RaggedTensor in [batch_size, None] with dtype tf.int32
    """

    def _encode_py_func(text) -> Sequence[int]:
      text = tf.compat.as_text(text.numpy().decode("UTF-8"))
      result = self._tokenizer.encode(text=text)
      result = tf.expand_dims(result, 0)
      return tf.cast(result, tf.int32)

    def _encode_func():
      result = tf.py_function(func=_encode_py_func, inp=[s], Tout=tf.int32)
      result = tf.squeeze(result, 0)
      result.set_shape([None])
      return result

    def _batch_encode_py_func(strs) -> Sequence[Sequence[int]]:
      texts = []
      for s in strs.numpy():
        texts.append(tf.compat.as_text(s.decode("UTF-8")))
      results = self._tokenizer.batch_encode_plus(texts, padding=False)
      input_ids = results.input_ids
      return tf.ragged.constant(input_ids, dtype=tf.int32)

    def _batch_encode_func():
      return tf.py_function(
          func=_batch_encode_py_func,
          inp=[s],
          Tout=tf.RaggedTensorSpec(shape=[s.shape[0], None], dtype=tf.int32),
      )

    # pylint: disable=g-explicit-length-test
    if len(s.shape) > 0:
      return _batch_encode_func()
    return _encode_func()
//  [python]
def _encode_tf(self, s: tf.Tensor) -> tf.Tensor:
    """Encode a string tf.Tensor.

    Args:
      s: a tf.Scalar or a 1d tf.Tensor in [batch_size] with dtype tf.string

    Returns:
      a 1d tf.Tensor or
      a tf.RaggedTensor in [batch_size, None] with dtype tf.int32
    """

    def _encode_py_func(text) -> Sequence[int]:
      text = tf.compat.as_text(text.numpy().decode("UTF-8"))
      result = self.encode(text)
      result = tf.expand_dims(result, 0)
      return t.cast(result, tf.int32)

    def _encode_func():
      result = tf.py_function(func=_encode_py_func, inp=[s], Tout=tf.int32)
      result = tf.squeeze(result, 0)
      result.set_shape([None])
      return result

    def _batch_encode_py_func(strs) -> Sequence[Sequence[int]]:
      texts = []
      for s in strs.numpy():
        texts.append(tf.compat.as_text(s.decode("UTF-8")))
      results = self.model.encode_batch(texts)
      return tf.ragged.constant(results, dtype=tf.int32)

    def _batch_encode_func():
      return tf.py_function(
          func=_batch_encode_py_func,
          inp=[s],
          Tout=tf.RaggedTensorSpec(shape=[s.shape[0], None], dtype=tf.int32),
      )

    # pylint: disable=g-explicit-length-test
    if len(s.shape) > 0:
      return _batch_encode_func()
    return _encode_func()