jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def call(self,
           inputs: tf.Tensor,
           image_info: Optional[tf.Tensor] = None) -> Dict[str, tf.Tensor]:
    """Generates panoptic segmentation masks."""
    # (batch_size, num_rois, 4) in absolute coordinates.
    detection_boxes = tf.to_float(inputs['detection_boxes'])
    # (batch_size, num_rois)
    detection_classes = tf.to_int32(inputs['detection_classes'])
    # (batch_size, num_rois)
    detection_scores = inputs['detection_scores']
    # (batch_size, num_rois, mask_height, mask_width)
    detections_masks = inputs['detection_masks']
    # (batch_size, height, width, num_semantic_classes)
    segmentation_outputs = inputs['segmentation_outputs']

    if self._rescale_predictions:
      # (batch_size, 2)
      original_size = tf.to_float(image_info[:, 0, :])
      desired_size = tf.to_float(image_info[:, 1, :])
      image_scale = tf.to_float(image_info[:, 2, :])
      offset = tf.to_float(image_info[:, 3, :])
      rescale_size = tf.math.ceil(desired_size / image_scale)
      # (batch_size, output_height, output_width, num_semantic_classes)
      segmentation_outputs = (
          spatial_transform_ops.bilinear_resize_with_crop_and_pad(
              segmentation_outputs,
              rescale_size,
              crop_offset=offset,
              crop_size=original_size,
              output_size=self._output_size))
      # (batch_size, 1, 4)
      image_scale = tf.tile(image_scale, multiples=[1, 2])[:, tf.newaxis]
      detection_boxes /= image_scale
    else:
      # (batch_size, output_height, output_width, num_semantic_classes)
      segmentation_outputs = tf.image.resize(
          segmentation_outputs, size=self._output_size, method='bilinear')

    # (batch_size, output_height, output_width)
    instance_mask, instance_category_mask = self._generate_instances(
        detection_boxes, detection_classes, detection_scores, detections_masks)

    # (batch_size, output_height, output_width)
    stuff_category_mask = self._generate_stuffs(segmentation_outputs)

    # (batch_size, output_height, output_width)
    category_mask = tf.where((stuff_category_mask != self._void_class_label) &
                             (instance_category_mask == self._void_class_label),
                             stuff_category_mask + self._stuff_classes_offset,
                             instance_category_mask)

    return {'instance_mask': instance_mask, 'category_mask': category_mask}
//  [python]
def call(self,
           inputs: tf.Tensor,
           image_info: Optional[tf.Tensor] = None) -> Dict[str, tf.Tensor]:
    """Generates panoptic segmentation masks."""
    # (batch_size, num_rois, 4) in absolute coordinates.
    detection_boxes = tf.cast(inputs['detection_boxes'], tf.float32)
    # (batch_size, num_rois)
    detection_classes = tf.cast(inputs['detection_classes'], tf.int32)
    # (batch_size, num_rois)
    detection_scores = inputs['detection_scores']
    # (batch_size, num_rois, mask_height, mask_width)
    detections_masks = inputs['detection_masks']
    # (batch_size, height, width, num_semantic_classes)
    segmentation_outputs = inputs['segmentation_outputs']

    if self._rescale_predictions:
      # (batch_size, 2)
      original_size = tf.cast(image_info[:, 0, :], tf.float32)
      desired_size = tf.cast(image_info[:, 1, :], tf.float32)
      image_scale = tf.cast(image_info[:, 2, :], tf.float32)
      offset = tf.cast(image_info[:, 3, :], tf.float32)
      rescale_size = tf.math.ceil(desired_size / image_scale)
      # (batch_size, output_height, output_width, num_semantic_classes)
      segmentation_outputs = (
          spatial_transform_ops.bilinear_resize_with_crop_and_pad(
              segmentation_outputs,
              rescale_size,
              crop_offset=offset,
              crop_size=original_size,
              output_size=self._output_size))
      # (batch_size, 1, 4)
      image_scale = tf.tile(image_scale, multiples=[1, 2])[:, tf.newaxis]
      detection_boxes /= image_scale
    else:
      # (batch_size, output_height, output_width, num_semantic_classes)
      segmentation_outputs = tf.image.resize(
          segmentation_outputs, size=self._output_size, method='bilinear')

    # (batch_size, output_height, output_width)
    instance_mask, instance_category_mask = self._generate_instances(
        detection_boxes, detection_classes, detection_scores, detections_masks)

    # (batch_size, output_height, output_width)
    stuff_category_mask = self._generate_stuffs(segmentation_outputs)

    # (batch_size, output_height, output_width)
    category_mask = tf.where((stuff_category_mask != self._void_class_label) &
                             (instance_category_mask == self._void_class_label),
                             stuff_category_mask + self._stuff_classes_offset,
                             instance_category_mask)

    return {'instance_mask': instance_mask, 'category_mask': category_mask}
//  [python]
def testTrainPerArmAgent(self):
    obs_spec = bandit_spec_utils.create_per_arm_observation_spec(2, 3, 3)
    time_step_spec = ts.time_step_spec(obs_spec)
    reward_net = (
        global_and_arm_feature_network.create_feed_forward_common_tower_network(
            obs_spec, (4, 3), (3, 4), (4, 2)
        )
    )
    optimizer = tf.compat.v1.train.GradientDescentOptimizer(learning_rate=0.1)
    agent = neural_boltzmann_agent.NeuralBoltzmannAgent(
        time_step_spec,
        self._action_spec,
        reward_network=reward_net,
        optimizer=optimizer,
        temperature=0.1,
        accepts_per_arm_features=True,
    )
    observations = {
        bandit_spec_utils.GLOBAL_FEATURE_KEY: tf.constant(
            [[1, 2], [3, 4]], dtype=tf.float32
        ),
        bandit_spec_utils.PER_ARM_FEATURE_KEY: tf.cast(
            tf.reshape(tf.range(18), shape=[2, 3, 3]), dtype=tf.float32
        ),
    }
    time_steps = ts.restart(observations, batch_size=2)
    policy = agent.policy
    action_step = policy.action(time_steps)
    self.evaluate(tf.compat.v1.initialize_all_variables())
    actions = self.evaluate(action_step.action)
    self.assertAllEqual(actions.shape, (2,))
//  [python]
def attention_mask(nd, ns, *, dtype):
    """1's in the lower triangle, counting from the lower right corner.

    Same as tf.matrix_band_part(tf.ones([nd, ns]), -1, ns-nd), but doesn't produce garbage on TPUs.
    """
    i = tf.range(nd)[:,None]
    j = tf.range(ns)
    m = i >= j - ns + nd
    return t.cast(m, dtype)
//  [python]
def attention_mask(nd, ns, *, dtype):
    """1's in the lower triangle, counting from the lower right corner.

    Same as tf.matrix_band_part(tf.ones([nd, ns]), -1, ns-nd), but doesn't produce garbage on TPUs.
    """
    i = tf.range(nd)[:,None]
    j = tf.range(ns)
    m = i >= j - ns + nd
    return tf.cast(m, dtype)
//  [python]
def compute_num_points_in_box_3d(point, box, name=None):
  """Computes the number of points in each box given a set of points and boxes.

  Args:
    point: [N, 3] tensor. Inner dims are: [x, y, z].
    box: [M, 7] tenosr. Inner dims are: [center_x, center_y, center_z, length,
      width, height, heading].
    name: tf name scope.

  Returns:
    num_points_in_box: [M] int32 tensor.
  """

  with tf.compat.v1.name_scope(name, 'ComputeNumPointsInBox3D', [point, box]):
    # [N, M]
    point_in_box = tf.to_int32(is_within_box_3d(point, box, name))
    num_points_in_box = tf.reduce_sum(input_tensor=point_in_box, axis=0)
    return num_points_in_box
//  [python]
def compute_num_points_in_box_3d(point, box, name=None):
  """Computes the number of points in each box given a set of points and boxes.

  Args:
    point: [N, 3] tensor. Inner dims are: [x, y, z].
    box: [M, 7] tenosr. Inner dims are: [center_x, center_y, center_z, length,
      width, height, heading].
    name: tf name scope.

  Returns:
    num_points_in_box: [M] int32 tensor.
  """

  with tf.compat.v1.name_scope(name, 'ComputeNumPointsInBox3D', [point, box]):
    # [N, M]
    point_in_box = tf.cast(is_within_box_3d(point, box, name), dtype=tf.int32)
    num_points_in_box = tf.reduce_sum(input_tensor=point_in_box, axis=0)
    return num_points_in_box
//  [python]
def _build_training_data(self):
        """Preprocess training dataset

        Return single batch of training dataset
        """

        def cast(x, y):
            return tf.cast(x, tf.float32), y

        x_raw = tf.random.uniform(
            minval=-0.5, maxval=0.5, shape=[self.training_set_size, self.num_features]
        )
        target_coeffs = tf.random.normal(
            shape=[self.num_features, 1], mean=3, stddev=2.0
        )
        y_raw = tf.matmul(x_raw, target_coeffs)

        return cast(x_raw, y_raw)
//  [python]
def is_within_box_2d(point, box):
  """Checks whether a point is in a BEV box given a set of points and boxes.

  Args:
    point: [N, 2] tensor. Inner dims are: [x, y].
    box: [M, 5] tensor. Inner dims are: [center_x, center_y, length, width,
      heading].

  Returns:
    point_in_box; [N, M] boolean tensor.
  """

  with tf.name_scope('IsWithinBox2D'):
    center = box[:, 0:2]
    dim = box[:, 2:4]
    heading = box[:, 4]
    # [M, 2, 2]
    rotation = transform_utils.get_yaw_rotation_2d(heading)
    # [M, 3, 3]
    transform = transform_utils.get_transform(rotation, center)
    # [M, 3, 3]
    transform = tf.linalg.inv(transform)
    # [M, 2, 2]
    rotation = transform[:, 0:2, 0:2]
    # [M, 2]
    translation = transform[:, 0:2, 2]

    # [N, M, 2]
    point_in_box_frame = tf.einsum('nj,mij->nmi', point, rotation) + translation
    # [N, M, 2]
    point_in_box = tf.logical_and(
        tf.logical_and(point_in_box_frame <= dim * 0.5,
                       point_in_box_frame >= -dim * 0.5),
        tf.reduce_all(tf.not_equal(dim, 0), axis=-1, keepdims=True))
    # [N, M]
    point_in_box = tf.reduce_prod(
        input_tensor=tf.to_int32(point_in_box), axis=-1,
    )
    point_in_box = tf.equal(point_in_box, 1)

    return point_in_box
//  [python]
def is_within_box_2d(point, box):
  """Checks whether a point is in a BEV box given a set of points and boxes.

  Args:
    point: [N, 2] tensor. Inner dims are: [x, y].
    box: [M, 5] tensor. Inner dims are: [center_x, center_y, length, width,
      heading].

  Returns:
    point_in_box; [N, M] boolean tensor.
  """

  with tf.name_scope('IsWithinBox2D'):
    center = box[:, 0:2]
    dim = box[:, 2:4]
    heading = box[:, 4]
    # [M, 2, 2]
    rotation = transform_utils.get_yaw_rotation_2d(heading)
    # [M, 3, 3]
    transform = transform_utils.get_transform(rotation, center)
    # [M, 3, 3]
    transform = tf.linalg.inv(transform)
    # [M, 2, 2]
    rotation = transform[:, 0:2, 0:2]
    # [M, 2]
    translation = transform[:, 0:2, 2]

    # [N, M, 2]
    point_in_box_frame = tf.einsum('nj,mij->nmi', point, rotation) + translation
    # [N, M, 2]
    point_in_box = tf.logical_and(
        tf.logical_and(point_in_box_frame <= dim * 0.5,
                       point_in_box_frame >= -dim * 0.5),
        tf.reduce_all(tf.not_equal(dim, 0), axis=-1, keepdims=True))
    # [N, M]
    point_in_box = tf.cast(
        tf.reduce_prod(
            input_tensor=tf.cast(point_in_box, dtype=tf.int32), axis=-1),
        dtype=tf.bool)

    return point_in_box
//  [python]
def call(self, inputs, mask):
    mask = t.cast(mask, inputs.dtype)
    return mask
//  [python]
def call(self, inputs, mask):
    mask = tf.cast(mask, inputs.dtype)
    return mask
//  [python]
def _resample_feature_map(self,
                            inputs,
                            input_level,
                            target_level,
                            target_num_filters=256):
    x = inputs
    _, _, _, input_num_filters = x.get_shape().as_list()
    if input_num_filters != target_num_filters:
      x = self._conv_op(
          filters=target_num_filters,
          kernel_size=1,
          padding='same',
          **self._conv_kwargs)(x)
      x = self._norm_op(**self._norm_kwargs)(x)

    if input_level < target_level:
      stride = int(2 ** (target_level - input_level))
      return tf_keras.layers.MaxPool2D(
          pool_size=stride, strides=stride, padding='same')(x)
    if input_level > target_level:
      scale = int(2 ** (input_level - target_level))
      return spatial_transform_ops.nearest_upsampling(x, scale=scale)

    # Force output x to be the same dtype as mixed precision policy. This avoids
    # dtype mismatch when one input (by default float32 dtype) does not meet all
    # the above conditions and is output unchanged, while other inputs are
    # processed to have different dtype, e.g., using bfloat16 on TPU.
    compute_dtype = tf_keras.layers.Layer().dtype_policy.compute_dtype
    if (compute_dtype is not None) and (x.dtype != compute_dtype):
      return tf.cast(x, dtype=compute_dtype)
    else:
      return x
//  [python]
def call(self, token_ids, seq_length):
    if self.parameters.mode in [base_layers.PREDICT, base_layers.TFLITE]:
      mask_rank2 = tf.ones(tf.shape(token_ids), dtype=tf.int32)
      seq_length = tf.reduce_sum(mask_rank2, axis=1)
      pos_indices = tf.cumsum(mask_rank2, axis=1, exclusive=True)
      pos_indices = tf.to_int32(pos_indices)
      pos_indices = tf.reshape(pos_indices, [1, -1])
    else:
      mask_rank2 = tf.sequence_mask(
          seq_length, tf.shape(token_ids)[1], dtype=tf.float32)
      pos_indices = tf.cumsum(mask_rank2, axis=1, exclusive=True)
      pos_indices = tf.to_int32(pos_indices)

    input_values = self.embedding(token_ids)
    pos_values = self.positional_embedding(pos_indices)
    input_embeds = self.qact(self.ln(input_values + pos_values))

    if self.bottleneck_layer is not None:
      maskr3 = tf.expand_dims(mask_rank2, axis=2)
      maskr3 = tf.to_float(maskr3)
      bottleneck_output = self.bottleneck_layer(input_embeds, maskr3)
    else:
      bottleneck_output = input_embeds

    gbst_output = self.gbst_layer(bottleneck_output, seq_length)
    if self.parameters.mode in [base_layers.PREDICT, base_layers.TFLITE]:
      mask_rank2 = tf.ones(tf.shape(gbst_output)[:-1], dtype=tf.float32)
      seq_length = tf.reduce_sum(mask_rank2, axis=1)
    else:
      seq_length = seq_length / self.gbst_downsample_rate

    if self.pool_windows:
      outputs, mask = self.transformer_encoder_layer(gbst_output,
                                                     seq_length)
      inverse_normalizer = tf.math.reciprocal(tf.reduce_sum(mask))

      pre_logits = self.attention_pool(outputs, mask, inverse_normalizer)
    else:
      outputs = self.transformer_encoder_layer(gbst_output, seq_length)
      mask = tf.sequence_mask(
          seq_length, tf.shape(outputs)[1], dtype=tf.float32)
      inverse_normalizer = tf.math.reciprocal(tf.reduce_sum(mask))
      maskr3 = tf.expand_dims(mask, axis=2)
      pre_logits = self.attention_pool(outputs, maskr3, inverse_normalizer)
    if self.num_classes:
      return self.final_fc(pre_logits)
    else:
      return pre_logits
//  [python]
def call(self, token_ids, seq_length):
    if self.parameters.mode in [base_layers.PREDICT, base_layers.TFLITE]:
      mask_rank2 = tf.ones(tf.shape(token_ids), dtype=tf.int32)
      seq_length = tf.reduce_sum(mask_rank2, axis=1)
      pos_indices = tf.cumsum(mask_rank2, axis=1, exclusive=True)
      pos_indices = tf.cast(pos_indices, dtype=tf.int32)
      pos_indices = tf.reshape(pos_indices, [1, -1])
    else:
      mask_rank2 = tf.sequence_mask(
          seq_length, tf.shape(token_ids)[1], dtype=tf.float32)
      pos_indices = tf.cumsum(mask_rank2, axis=1, exclusive=True)
      pos_indices = tf.cast(pos_indices, dtype=tf.int32)

    input_values = self.embedding(token_ids)
    pos_values = self.positional_embedding(pos_indices)
    input_embeds = self.qact(self.ln(input_values + pos_values))

    if self.bottleneck_layer is not None:
      maskr3 = tf.expand_dims(mask_rank2, axis=2)
      maskr3 = tf.cast(maskr3, tf.float32)
      bottleneck_output = self.bottleneck_layer(input_embeds, maskr3)
    else:
      bottleneck_output = input_embeds

    gbst_output = self.gbst_layer(bottleneck_output, seq_length)
    if self.parameters.mode in [base_layers.PREDICT, base_layers.TFLITE]:
      mask_rank2 = tf.ones(tf.shape(gbst_output)[:-1], dtype=tf.float32)
      seq_length = tf.reduce_sum(mask_rank2, axis=1)
    else:
      seq_length = seq_length / self.gbst_downsample_rate

    if self.pool_windows:
      outputs, mask = self.transformer_encoder_layer(gbst_output,
                                                     seq_length)
      inverse_normalizer = tf.math.reciprocal(tf.reduce_sum(mask))

      pre_logits = self.attention_pool(outputs, mask, inverse_normalizer)
    else:
      outputs = self.transformer_encoder_layer(gbst_output, seq_length)
      mask = tf.sequence_mask(
          seq_length, tf.shape(outputs)[1], dtype=tf.float32)
      inverse_normalizer = tf.math.reciprocal(tf.reduce_sum(mask))
      maskr3 = tf.expand_dims(mask, axis=2)
      pre_logits = self.attention_pool(outputs, maskr3, inverse_normalizer)
    if self.num_classes:
      return self.final_fc(pre_logits)
    else:
      return pre_logits
//  [python]
def _build_testing_data(self):
        """Preprocess testing dataset

        Return single batch of testing dataset
        """

        def cast(x, y):
            return tf.cast(x, tf.float32), tf.expand_dims(y, 0)

        x_raw = tf.random.uniform(
            minval=-0.5, maxval=0.5, shape=[self.training_set_size, self.num_features]
        )
        target_coeffs = tf.random.normal(shape=(self.num_features), mean=3, stddev=2.0)
        y_raw = tf.matmul(x_raw, target_coeffs)

        return cast(x_raw, y_raw)
//  [python]
def invert_permutation(
    x: Union[tf.Tensor, tf.Variable, list, tuple],
    /,
) -> Union[tf.Tensor, tf.Variable]:
    return tf.cast(tf.math.invert_permutation(x), tf.int64)
//  [python]
def atomic_dress(tensors, dress, dtype=tf.float32):
    """Assign an energy to each specified elems

    Args:
        dress (dict): dictionary consisting the atomic energies
    """
    elem = tensors['elems']
    e_dress = tf.zeros_like(elem, dtype)
    for k, val in dress.items():
        indices = t.cast(tf.equal(elem, k), dtype)
        e_dress += indices * t.cast(val, dtype)
    n_batch = tf.reduce_max(tensors['ind_1'])+1
    e_dress = tf.math.unsorted_segment_sum(
        e_dress, tensors['ind_1'][:, 0], n_batch)
    return e_dress
//  [python]
def atomic_dress(tensors, dress, dtype=tf.float32):
    """Assign an energy to each specified elems

    Args:
        dress (dict): dictionary consisting the atomic energies
    """
    elem = tensors['elems']
    e_dress = tf.zeros_like(elem, dtype)
    for k, val in dress.items():
        indices = tf.cast(tf.equal(elem, k), dtype)
        e_dress += indices * tf.cast(val, dtype)
    n_batch = tf.reduce_max(tensors['ind_1'])+1
    e_dress = tf.math.unsorted_segment_sum(
        e_dress, tensors['ind_1'][:, 0], n_batch)
    return e_dress
//  [python]
def _generalized_kernel(x, y, is_query, projection_matrix, f, h):
  """Generalized kernel in RETHINKING ATTENTION WITH PERFORMERS.

  Args:
    x: The feature being transformed with shape [B, T, N ,H].
    y: The extra stats-tensor of shape [B, T, N ,H].
    is_query: True if x is a query-tensor.
    projection_matrix: The matrix with shape [M, H] that we projecct x to, where
      M is the number of projections.
    f: A non-linear function applied on x or projected x.
    h: A muliplier which is a function of x applied after projected and
      transformed. Only applied if projection_matrix is not None.

  Returns:
    Transformed feature.
  """
  del y
  del is_query
  if projection_matrix is None:
    return h(x) * f(x)
  else:
    x_projected = tf.einsum("BTNH,MH->BTNM", x, projection_matrix)
    return h(x) * f(x_projected) / tf.math.sqrt(
        tf.cast(tf.shape(projection_matrix)[0], tf.float32))
//  [python]
def attn(x, scope, n_state, *, past, hparams):
    assert x.shape.ndims == 3  # Should be [batch, sequence, features]
    assert n_state % hparams.n_head == 0
    if past is not None:
        assert past.shape.ndims == 5  # Should be [batch, 2, heads, sequence, features], where 2 is [k, v]

    def split_heads(x):
        # From [batch, sequence, features] to [batch, heads, sequence, features]
        return tf.transpose(split_states(x, hparams.n_head), [0, 2, 1, 3])

    def merge_heads(x):
        # Reverse of split_heads
        return merge_states(tf.transpose(x, [0, 2, 1, 3]))

    def mask_attn_weights(w):
        # w has shape [batch, heads, dst_sequence, src_sequence], where information flows from src to dst.
        _, _, nd, ns = shape_list(w)
        b = attention_mask(nd, ns, dtype=w.dtype)
        b = tf.reshape(b, [1, 1, nd, ns])
        w = w*b - tf.cast(1e10, w.dtype)*(1-b)
        return w

    def multihead_attn(q, k, v):
        # q, k, v have shape [batch, heads, sequence, features]
        w = tf.matmul(q, k, transpose_b=True)
        w = w * tf.rsqrt(tf.cast(v.shape[-1].value, w.dtype))

        w = mask_attn_weights(w)
        w = softmax(w)
        a = tf.matmul(w, v)
        return a

    with tf.variable_scope(scope):
        c = conv1d(x, 'c_attn', n_state*3)
        q, k, v = map(split_heads, tf.split(c, 3, axis=2))
        present = tf.stack([k, v], axis=1)
        if past is not None:
            pk, pv = tf.unstack(past, axis=1)
            k = tf.concat([pk, k], axis=-2)
            v = tf.concat([pv, v], axis=-2)
        a = multihead_attn(q, k, v)
        a = merge_heads(a)
        a = conv1d(a, 'c_proj', n_state)
        return a, present
//  [python]
def get_atomic_dress(dataset, elems, key='e_data'):
    """Fit the atomic energy with a element dependent atomic dress

    Args:
        dataset: dataset to fit
        elems: a list of element numbers
        key: key of the property to fit

    Returns:
        atomic_dress: a dictionary comprising the atomic energy of each element
        error: residue error of the atomic dress
    """
    from pinn.io import sparse_batch

    def count_elems(tensors):
        tensors = tensors.copy()
        count = tf.equal(tf.expand_dims(
            tensors['elems'], 1), tf.expand_dims(elems, 0))
        count = tf.cast(count, tf.int32)
        count = tf.math.segment_sum(count, tensors['ind_1'][:, 0])
        return count, tensors[key]

    if 'ind_1' not in next(iter(dataset)):
        dataset = dataset.apply(sparse_batch(1))
    x, y = [], []
    for x_i, y_i in dataset.map(count_elems).as_numpy_iterator():
        x.append(x_i)
        y.append(y_i)

    x, y = np.concatenate(x, 0), np.concatenate(y, 0)
    beta = np.dot(np.dot(np.linalg.pinv(np.dot(x.T, x)), x.T), np.array(y))
    dress = {e: float(beta[i]) for (i, e) in enumerate(elems)}
    error = np.dot(x, beta) - y
    return dress, error
//  [python]
def lexsort(
    keys: Union[tf.Tensor, tf.Variable],
    /,
    *,
    axis: int = -1,
    out: Optional[Union[tf.Tensor, tf.Variable]] = None,
) -> Union[tf.Tensor, tf.Variable]:
    shape = keys.shape.as_list()
    if len(shape) == 1:
        return t.cast(tf.argsort(keys, axis=axis, stable=True), tf.int64)
    if shape[0] == 0:
        raise TypeError("need sequence of keys with len > 0 in lexsort")
    if len(shape) == 2 and shape[1] == 1:
        return t.cast(tf.convert_to_tensor([0]), tf.int64)
    result = tf.argsort(keys[0], axis=axis, stable=True)
    if shape[0] == 1:
        return t.cast(result, tf.int64)
    for i in range(1, shape[0]):
        key = keys[i]
        ind = tf.gather(key, result)
        temp = tf.argsort(ind, axis=axis, stable=True)
        result = tf.gather(result, temp)
    return t.cast(result, tf.int64)
//  [python]
def lexsort(
    keys: Union[tf.Tensor, tf.Variable],
    /,
    *,
    axis: int = -1,
    out: Optional[Union[tf.Tensor, tf.Variable]] = None,
) -> Union[tf.Tensor, tf.Variable]:
    shape = keys.shape.as_list()
    if len(shape) == 1:
        return tf.cast(tf.argsort(keys, axis=axis, stable=True), tf.int64)
    if shape[0] == 0:
        raise TypeError("need sequence of keys with len > 0 in lexsort")
    if len(shape) == 2 and shape[1] == 1:
        return tf.cast(tf.convert_to_tensor([0]), tf.int64)
    result = tf.argsort(keys[0], axis=axis, stable=True)
    if shape[0] == 1:
        return tf.cast(result, tf.int64)
    for i in range(1, shape[0]):
        key = keys[i]
        ind = tf.gather(key, result)
        temp = tf.argsort(ind, axis=axis, stable=True)
        result = tf.gather(result, temp)
    return tf.cast(result, tf.int64)
//  [python]
def __init__(self, max_seq, embedding_dim, **kwargs):
        super(PositionEmbeddingV2, self).__init__(**kwargs)
        angle_rads = PositionEmbeddingV2.__get_angles(np.arange(max_seq)[:, np.newaxis],
                                                      np.arange(embedding_dim)[np.newaxis, :],
                                                      embedding_dim)

        # apply sin to even indices in the array; 2i
        sines = np.sin(angle_rads[:, 0::2])

        # apply cos to odd indices in the array; 2i+1
        cosines = np.cos(angle_rads[:, 1::2])

        pos_encoding = np.concatenate([sines, cosines], axis=-1)

        pos_encoding = pos_encoding[np.newaxis, ...]

        self.sinusoid = t.cast(pos_encoding, dtype=t.float32)
//  [python]
def __init__(self, max_seq, embedding_dim, **kwargs):
        super(PositionEmbeddingV2, self).__init__(**kwargs)
        angle_rads = PositionEmbeddingV2.__get_angles(np.arange(max_seq)[:, np.newaxis],
                                                      np.arange(embedding_dim)[np.newaxis, :],
                                                      embedding_dim)

        # apply sin to even indices in the array; 2i
        sines = np.sin(angle_rads[:, 0::2])

        # apply cos to odd indices in the array; 2i+1
        cosines = np.cos(angle_rads[:, 1::2])

        pos_encoding = np.concatenate([sines, cosines], axis=-1)

        pos_encoding = pos_encoding[np.newaxis, ...]

        self.sinusoid = tf.cast(pos_encoding, dtype=tf.float32)
//  [python]
def _qe_masking(qe):
        mask = tf.sequence_mask(
            tf.range(qe.shape[-1] -1, qe.shape[-1] - qe.shape[-2] -1, -1), qe.shape[-1])

        mask = tf.logical_not(mask)
        mask = tf.to_float(mask)

        return mask * qe
//  [python]
def _qe_masking(qe):
        mask = tf.sequence_mask(
            tf.range(qe.shape[-1] -1, qe.shape[-1] - qe.shape[-2] -1, -1), qe.shape[-1])

        mask = tf.logical_not(mask)
        mask = tf.cast(mask, tf.float32)

        return mask * qe
//  [python]
def call(self, x, mask=None, training=False):
        weights = []
        # adding embedding and position encoding.
        x = self.embedding(x)  # (batch_size, input_seq_len, d_model)
        x *= tf.sqrt(tf.to_float(self.d_model))
        x = self.pos_encoding(x)
        x = self.dropout(x, training=training)
        for i in range(self.num_layers):
            x, w = self.enc_layers[i](x, mask, training=training)
            weights.append(w)
        return x, weights
//  [python]
def call(self, x, mask=None, training=False):
        weights = []
        # adding embedding and position encoding.
        x = self.embedding(x)  # (batch_size, input_seq_len, d_model)
        x *= tf.math.sqrt(tf.cast(self.d_model, tf.float32))
        x = self.pos_encoding(x)
        x = self.dropout(x, training=training)
        for i in range(self.num_layers):
            x, w = self.enc_layers[i](x, mask, training=training)
            weights.append(w)
        return x, weights
//  [python]
def predict(self, inputs):
        """Returns the logits for the test items.

        Args:
            inputs (tf.Tensor): Input tensor.

        Returns:
             tf.Tensor: Output tensor.
        """
        training = False
        input_seq = inputs["input_seq"]
        candidate = inputs["candidate"]

        mask = tf.expand_dims(tf.to_float(tf.not_equal(input_seq, 0)), -1)
        seq_embeddings, positional_embeddings = self.embedding(input_seq)
        seq_embeddings += positional_embeddings
        # seq_embeddings = self.dropout_layer(seq_embeddings)
        seq_embeddings *= mask
        seq_attention = seq_embeddings
        seq_attention = self.encoder(seq_attention, training, mask)
        seq_attention = self.layer_normalization(seq_attention)  # (b, s, d)
        seq_emb = tf.reshape(
            seq_attention,
            [tf.shape(input_seq)[0] * self.seq_max_len, self.embedding_dim],
        )  # (b*s, d)
        candidate_emb = self.item_embedding_layer(candidate)  # (b, s, d)
        candidate_emb = tf.transpose(candidate_emb, perm=[0, 2, 1])  # (b, d, s)

        test_logits = tf.matmul(seq_emb, candidate_emb)
        # (200, 100) * (1, 101, 100)'

        test_logits = tf.reshape(
            test_logits,
            [tf.shape(input_seq)[0], self.seq_max_len, 1 + self.num_neg_test],
        )  # (1, 200, 101)
        test_logits = test_logits[:, -1, :]  # (1, 101)
        return test_logits
//  [python]
def predict(self, inputs):
        """Returns the logits for the test items.

        Args:
            inputs (tf.Tensor): Input tensor.

        Returns:
             tf.Tensor: Output tensor.
        """
        training = False
        input_seq = inputs["input_seq"]
        candidate = inputs["candidate"]

        mask = tf.expand_dims(tf.cast(tf.not_equal(input_seq, 0), tf.float32), -1)
        seq_embeddings, positional_embeddings = self.embedding(input_seq)
        seq_embeddings += positional_embeddings
        # seq_embeddings = self.dropout_layer(seq_embeddings)
        seq_embeddings *= mask
        seq_attention = seq_embeddings
        seq_attention = self.encoder(seq_attention, training, mask)
        seq_attention = self.layer_normalization(seq_attention)  # (b, s, d)
        seq_emb = tf.reshape(
            seq_attention,
            [tf.shape(input_seq)[0] * self.seq_max_len, self.embedding_dim],
        )  # (b*s, d)
        candidate_emb = self.item_embedding_layer(candidate)  # (b, s, d)
        candidate_emb = tf.transpose(candidate_emb, perm=[0, 2, 1])  # (b, d, s)

        test_logits = tf.matmul(seq_emb, candidate_emb)
        # (200, 100) * (1, 101, 100)'

        test_logits = tf.reshape(
            test_logits,
            [tf.shape(input_seq)[0], self.seq_max_len, 1 + self.num_neg_test],
        )  # (1, 200, 101)
        test_logits = test_logits[:, -1, :]  # (1, 101)
        return test_logits
//  [python]
def call(self, x, training):
        """Model forward pass.

        Args:
            x (tf.Tensor): Input tensor.
            training (tf.Tensor): Training tensor.

        Returns:
            tf.Tensor, tf.Tensor, tf.Tensor:
            - Logits of the positive examples.
            - Logits of the negative examples.
            - Mask for nonzero targets
        """

        input_seq = x["input_seq"]
        pos = x["positive"]
        neg = x["negative"]

        mask = tf.expand_dims(tf.to_float(tf.not_equal(input_seq, 0)), -1)
        seq_embeddings, positional_embeddings = self.embedding(input_seq)

        # add positional embeddings
        seq_embeddings += positional_embeddings

        # dropout
        seq_embeddings = self.dropout_layer(seq_embeddings)

        # masking
        seq_embeddings *= mask

        # --- ATTENTION BLOCKS ---
        seq_attention = seq_embeddings
        seq_attention = self.encoder(seq_attention, training, mask)
        seq_attention = self.layer_normalization(seq_attention)  # (b, s, d)

        # --- PREDICTION LAYER ---
        # user's sequence embedding
        pos = self.mask_layer(pos)
        neg = self.mask_layer(neg)

        pos = tf.reshape(pos, [tf.shape(input_seq)[0] * self.seq_max_len])
        neg = tf.reshape(neg, [tf.shape(input_seq)[0] * self.seq_max_len])
        pos_emb = self.item_embedding_layer(pos)
        neg_emb = self.item_embedding_layer(neg)
        seq_emb = tf.reshape(
            seq_attention,
            [tf.shape(input_seq)[0] * self.seq_max_len, self.embedding_dim],
        )  # (b*s, d)

        pos_logits = tf.reduce_sum(pos_emb * seq_emb, -1)
        neg_logits = tf.reduce_sum(neg_emb * seq_emb, -1)

        pos_logits = tf.expand_dims(pos_logits, axis=-1)  # (bs, 1)
        # pos_prob = tf.keras.layers.Dense(1, activation='sigmoid')(pos_logits)  # (bs, 1)

        neg_logits = tf.expand_dims(neg_logits, axis=-1)  # (bs, 1)
        # neg_prob = tf.keras.layers.Dense(1, activation='sigmoid')(neg_logits)  # (bs, 1)

        # output = tf.concat([pos_logits, neg_logits], axis=0)

        # masking for loss calculation
        istarget = tf.reshape(
            tf.to_float(tf.not_equal(pos, 0)),
            [tf.shape(input_seq)[0] * self.seq_max_len],
        )

        return pos_logits, neg_logits, istarget
//  [python]
def call(self, x, training):
        """Model forward pass.

        Args:
            x (tf.Tensor): Input tensor.
            training (tf.Tensor): Training tensor.

        Returns:
            tf.Tensor, tf.Tensor, tf.Tensor:
            - Logits of the positive examples.
            - Logits of the negative examples.
            - Mask for nonzero targets
        """

        input_seq = x["input_seq"]
        pos = x["positive"]
        neg = x["negative"]

        mask = tf.expand_dims(tf.cast(tf.not_equal(input_seq, 0), tf.float32), -1)
        seq_embeddings, positional_embeddings = self.embedding(input_seq)

        # add positional embeddings
        seq_embeddings += positional_embeddings

        # dropout
        seq_embeddings = self.dropout_layer(seq_embeddings)

        # masking
        seq_embeddings *= mask

        # --- ATTENTION BLOCKS ---
        seq_attention = seq_embeddings
        seq_attention = self.encoder(seq_attention, training, mask)
        seq_attention = self.layer_normalization(seq_attention)  # (b, s, d)

        # --- PREDICTION LAYER ---
        # user's sequence embedding
        pos = self.mask_layer(pos)
        neg = self.mask_layer(neg)

        pos = tf.reshape(pos, [tf.shape(input_seq)[0] * self.seq_max_len])
        neg = tf.reshape(neg, [tf.shape(input_seq)[0] * self.seq_max_len])
        pos_emb = self.item_embedding_layer(pos)
        neg_emb = self.item_embedding_layer(neg)
        seq_emb = tf.reshape(
            seq_attention,
            [tf.shape(input_seq)[0] * self.seq_max_len, self.embedding_dim],
        )  # (b*s, d)

        pos_logits = tf.reduce_sum(pos_emb * seq_emb, -1)
        neg_logits = tf.reduce_sum(neg_emb * seq_emb, -1)

        pos_logits = tf.expand_dims(pos_logits, axis=-1)  # (bs, 1)
        # pos_prob = tf.keras.layers.Dense(1, activation='sigmoid')(pos_logits)  # (bs, 1)

        neg_logits = tf.expand_dims(neg_logits, axis=-1)  # (bs, 1)
        # neg_prob = tf.keras.layers.Dense(1, activation='sigmoid')(neg_logits)  # (bs, 1)

        # output = tf.concat([pos_logits, neg_logits], axis=0)

        # masking for loss calculation
        istarget = tf.reshape(
            tf.cast(tf.not_equal(pos, 0), dtype=tf.float32),
            [tf.shape(input_seq)[0] * self.seq_max_len],
        )

        return pos_logits, neg_logits, istarget
//  [python]
def call(self, inputs, mask=None, weight_out=False, **kwargs):
        """
        :param inputs: a list of tensors. i.e) [Q, K, V]
        :param mask: mask tensor
        :param weight_out: decide to get weather weight or not
        :param kwargs:
        :return: final tensor ( output of attention )
        """
        q = inputs[0]
        q = self.Wq(q)
        q = tf.reshape(q, (q.shape[0], q.shape[1], self.h, -1))
        q = tf.transpose(q, (0, 2, 1, 3))  # batch, h, seq, dh

        k = inputs[1]
        k = self.Wk(k)
        k = tf.reshape(k, (k.shape[0], k.shape[1], self.h, -1))
        k = tf.transpose(k, (0, 2, 1, 3))

        v = inputs[2]
        v = self.Wv(v)
        v = tf.reshape(v, (v.shape[0], v.shape[1], self.h, -1))
        v = tf.transpose(v, (0, 2, 1, 3))

        Kt = tf.transpose(k, [0, 1, 3, 2])
        QKt = tf.matmul(q, Kt)
        logits = QKt
        logits = logits / math.sqrt(self.dh)

        if mask is not None:
            logits += (tf.to_float(mask) * -1e9)

        attention_weights = tf.nn.softmax(logits, -1)
        attention = tf.matmul(attention_weights, v)

        out = tf.transpose(attention, (0, 2, 1, 3))
        out = tf.reshape(out, (out.shape[0], -1, self.d))

        out = self.fc(out)

        return out, attention_weights
//  [python]
def call(self, inputs, mask=None, weight_out=False, **kwargs):
        """
        :param inputs: a list of tensors. i.e) [Q, K, V]
        :param mask: mask tensor
        :param weight_out: decide to get weather weight or not
        :param kwargs:
        :return: final tensor ( output of attention )
        """
        q = inputs[0]
        q = self.Wq(q)
        q = tf.reshape(q, (q.shape[0], q.shape[1], self.h, -1))
        q = tf.transpose(q, (0, 2, 1, 3))  # batch, h, seq, dh

        k = inputs[1]
        k = self.Wk(k)
        k = tf.reshape(k, (k.shape[0], k.shape[1], self.h, -1))
        k = tf.transpose(k, (0, 2, 1, 3))

        v = inputs[2]
        v = self.Wv(v)
        v = tf.reshape(v, (v.shape[0], v.shape[1], self.h, -1))
        v = tf.transpose(v, (0, 2, 1, 3))

        Kt = tf.transpose(k, [0, 1, 3, 2])
        QKt = tf.matmul(q, Kt)
        logits = QKt
        logits = logits / math.sqrt(self.dh)

        if mask is not None:
            logits += (tf.cast(mask, tf.float32) * -1e9)

        attention_weights = tf.nn.softmax(logits, -1)
        attention = tf.matmul(attention_weights, v)

        out = tf.transpose(attention, (0, 2, 1, 3))
        out = tf.reshape(out, (out.shape[0], -1, self.d))

        out = self.fc(out)

        return out, attention_weights
//  [python]
def call(self, x, mask, lookup_mask, training, enc_output=None):
        weights = []
        # adding embedding and position encoding.
        x = self.embedding(x)  # (batch_size, input_seq_len, d_model)
        x *= tf.sqrt(tf.cast(self.d_model, tf.float32))
        x = self.pos_encoding(x)
        x = self.dropout(x, training=training)
        for i in range(self.num_layers):
            x, w1, w2 = \
                self.dec_layers[i](x, enc_output, lookup_mask=lookup_mask, mask=mask, training=training, w_out=True)
            weights.append((w1, w2))

        return x, weights
//  [python]
def call(self, x, mask, lookup_mask, training, enc_output=None):
        weights = []
        # adding embedding and position encoding.
        x = self.embedding(x)  # (batch_size, input_seq_len, d_model)
        x *= tf.math.sqrt(tf.cast(self.d_model, tf.float32))
        x = self.pos_encoding(x)
        x = self.dropout(x, training=training)
        for i in range(self.num_layers):
            x, w1, w2 = \
                self.dec_layers[i](x, enc_output, lookup_mask=lookup_mask, mask=mask, training=training, w_out=True)
            weights.append((w1, w2))

        return x, weights
//  [python]
def check_embeddings_within_bounds(tensor: tf.Tensor, embed_dim: int, tensor_name: str = "input_ids") -> None:
    """
    `tf.gather`, on which TF embedding layers are based, won't check positive out of bound indices on GPU, returning
    zeros instead. This function adds a check against that dangerous silent behavior.

    Args:
        tensor (`tf.Tensor`): The tensor of indices to check.
        embed_dim (`int`): The embedding dimension.
        tensor_name (`str`, *optional*): The name of the tensor to use in the error message.
    """
    tf.debugging.assert_less(
        tensor,
        tf.to_int32(embed_dim),
        message=(
            f"The maximum value of {tensor_name} ({tf.math.reduce_max(tensor)}) must be smaller than the embedding "
            f"layer's input dimension ({embed_dim}). The likely cause is some problem at tokenization time."
        ),
    )
//  [python]
def check_embeddings_within_bounds(tensor: tf.Tensor, embed_dim: int, tensor_name: str = "input_ids") -> None:
    """
    `tf.gather`, on which TF embedding layers are based, won't check positive out of bound indices on GPU, returning
    zeros instead. This function adds a check against that dangerous silent behavior.

    Args:
        tensor (`tf.Tensor`): The tensor of indices to check.
        embed_dim (`int`): The embedding dimension.
        tensor_name (`str`, *optional*): The name of the tensor to use in the error message.
    """
    tf.debugging.assert_less(
        tensor,
        tf.cast(embed_dim, dtype=tensor.dtype),
        message=(
            f"The maximum value of {tensor_name} ({tf.math.reduce_max(tensor)}) must be smaller than the embedding "
            f"layer's input dimension ({embed_dim}). The likely cause is some problem at tokenization time."
        ),
    )
//  [python]
def __call__(self, q, k, v, mask):
    """Applies scaled dot product attention with softmax normalization.

    Args:
      q: The queries to the attention layer.
      k: The keys to the attention layer.
      v: The values to the attention layer.
      mask: The mask applied to the input to the softmax.

    Returns:
      A Tuple of layer outputs and attention weights.
    """
    normalization_constant = tf.sqrt(tf.to_float(tf.shape(k)[-1]))
    q = tf.transpose(q, [1, 0, 2])
    k = tf.transpose(k, [1, 0, 2])
    v = tf.transpose(v, [1, 0, 2])
    mask = tf.transpose(mask, [1, 0, 2])

    attn = tf.einsum('fbd,tbd->fbt', q, k) / normalization_constant

    attn -= tf.reduce_max(
        attn + tf.math.log(mask + 1e-9), axis=2, keepdims=True)
    attn = mask * tf.exp(attn)
    attn = tf.math.divide(
        attn,
        tf.reduce_sum(attn, axis=2, keepdims=True) + 1e-9,
    )

    output = tf.einsum('fbt,tbd->fbd', attn, v)
    output = tf.transpose(output, [1, 0, 2])

    return output, attn
//  [python]
def __call__(self, q, k, v, mask):
    """Applies scaled dot product attention with softmax normalization.

    Args:
      q: The queries to the attention layer.
      k: The keys to the attention layer.
      v: The values to the attention layer.
      mask: The mask applied to the input to the softmax.

    Returns:
      A Tuple of layer outputs and attention weights.
    """
    normalization_constant = tf.sqrt(tf.cast(tf.shape(k)[-1], dtype='float32'))
    q = tf.transpose(q, [1, 0, 2])
    k = tf.transpose(k, [1, 0, 2])
    v = tf.transpose(v, [1, 0, 2])
    mask = tf.transpose(mask, [1, 0, 2])

    attn = tf.einsum('fbd,tbd->fbt', q, k) / normalization_constant

    attn -= tf.reduce_max(
        attn + tf.math.log(mask + 1e-9), axis=2, keepdims=True)
    attn = mask * tf.exp(attn)
    attn = tf.math.divide(
        attn,
        tf.reduce_sum(attn, axis=2, keepdims=True) + 1e-9,
    )

    output = tf.einsum('fbt,tbd->fbd', attn, v)
    output = tf.transpose(output, [1, 0, 2])

    return output, attn
//  [python]
def invert_attention_mask(encoder_attention_mask: tf.Tensor) -> tf.Tensor:
    """
    Invert an attention mask (e.g., switches 0. and 1.).

    Args:
        encoder_attention_mask (`torch.Tensor`): An attention mask.

    Returns:
        `tf.Tensor`: The inverted attention mask.
    """
    if not isinstance(encoder_attention_mask, tf.Tensor):
        encoder_attention_mask = tf.convert_to_tensor(encoder_attention_mask)  # Catches stray NumPy inputs
    if encoder_attention_mask.shape.rank == 3:
        encoder_extended_attention_mask = encoder_attention_mask[:, None, :, :]
    if encoder_attention_mask.shape.rank == 2:
        encoder_extended_attention_mask = encoder_attention_mask[:, None, None, :]
    # T5 has a mask that can compare sequence ids, we can simulate this here with this transposition
    # Cf. https://github.com/tensorflow/mesh/blob/8d2465e9bc93129b913b5ccc6a59aa97abd96ec6/mesh_tensorflow
    # /transformer/transformer_layers.py#L270
    # encoder_extended_attention_mask = (encoder_extended_attention_mask ==
    # encoder_extended_attention_mask.transpose(-1, -2))
    encoder_extended_attention_mask = (
        tf.to_float(1) - encoder_extended_attention_mask
    ) * encoder_extended_attention_mask.dtype.min

    return encoder_extended_attention_mask
//  [python]
def invert_attention_mask(encoder_attention_mask: tf.Tensor) -> tf.Tensor:
    """
    Invert an attention mask (e.g., switches 0. and 1.).

    Args:
        encoder_attention_mask (`torch.Tensor`): An attention mask.

    Returns:
        `tf.Tensor`: The inverted attention mask.
    """
    if not isinstance(encoder_attention_mask, tf.Tensor):
        encoder_attention_mask = tf.convert_to_tensor(encoder_attention_mask)  # Catches stray NumPy inputs
    if encoder_attention_mask.shape.rank == 3:
        encoder_extended_attention_mask = encoder_attention_mask[:, None, :, :]
    if encoder_attention_mask.shape.rank == 2:
        encoder_extended_attention_mask = encoder_attention_mask[:, None, None, :]
    # T5 has a mask that can compare sequence ids, we can simulate this here with this transposition
    # Cf. https://github.com/tensorflow/mesh/blob/8d2465e9bc93129b913b5ccc6a59aa97abd96ec6/mesh_tensorflow
    # /transformer/transformer_layers.py#L270
    # encoder_extended_attention_mask = (encoder_extended_attention_mask ==
    # encoder_extended_attention_mask.transpose(-1, -2))
    encoder_extended_attention_mask = (
        tf.cast(1, encoder_attention_mask.dtype) - encoder_extended_attention_mask
    ) * encoder_extended_attention_mask.dtype.min

    return encoder_extended_attention_mask
//  [python]
def random_greedy_decode(prob_greedy, batch_size):
    noise = tf.less(tf.random.uniform(shape=[batch_size], dtype=tf.float32), prob_greedy)
    correct_answers = tf.zeros(batch_size, dtype=tf.int32)
    # prediction for each batch element for this timestep
    def apply_fn(unary_scores, mask):
        labels_mask_casted = t.cast(mask, unary_scores.dtype)
        masked_unary_scores = labels_mask_casted * unary_scores - (1 - labels_mask_casted) * 50
        greedy_answers = t.cast(tf.argmax(masked_unary_scores, axis=-1), tf.int32)
        return tf.where(noise,
                        greedy_answers,
                        correct_answers)
    return apply_fn
//  [python]
def random_greedy_decode(prob_greedy, batch_size):
    noise = tf.less(tf.random.uniform(shape=[batch_size], dtype=tf.float32), prob_greedy)
    correct_answers = tf.zeros(batch_size, dtype=tf.int32)
    # prediction for each batch element for this timestep
    def apply_fn(unary_scores, mask):
        labels_mask_casted = tf.cast(mask, unary_scores.dtype)
        masked_unary_scores = labels_mask_casted * unary_scores - (1 - labels_mask_casted) * 50
        greedy_answers = tf.cast(tf.argmax(masked_unary_scores, axis=-1), tf.int32)
        return tf.where(noise,
                        greedy_answers,
                        correct_answers)
    return apply_fn
//  [python]
def build_embed(inputs, features, index2words, keep_prob, is_training, create_embedding_lookup=True):
    embeddings = []
    word_embeddings = None
    for idx, (values, feature, index2word) in enumerate(zip(inputs, features, index2words)):
        if requires_vocab(feature):
            with tf.variable_scope("embedding_%d" % (idx,)):
                if create_embedding_lookup:
                    W, embedding = embedding_lookup(values,
                                                    dim=feature["dimension"],
                                                    size=len(index2word),
                                                    dtype=tf.float32,
                                                    mask_negative=True)
                    if feature["type"] == "word":
                        word_embeddings = (W, embedding, values)
                else:
                    embedding = tf.placeholder(tf.float32,
                                               feature_shape(feature) + [feature["dimension"]])
                tf.add_to_collection(EMBEDDED_INDICES, embedding)
                embeddings.append(embedding)
        elif requires_character_convolution(feature):
            embeddings.append(character_convolution(values, feature))
        elif feature["type"] == "bio":
            # b if different from previous tag
            start_is_b = tf.greater(values[:1], -1)
            b = tf.logical_and(tf.greater(values, -1), tf.concat([start_is_b, tf.not_equal(values[:-1], values[1:])], axis=0))
            i = tf.logical_and(tf.greater(values, -1), tf.concat([tf.zeros_like(start_is_b), tf.equal(values[:-1], values[1:])], axis=0))
            # O-tag is -1
            o = tf.less(values, 0)
            embeddings.append(tf.to_float(tf.stack([b, i, o], axis=-1, name="bool_BIO"), name="BIO"))
        else:
            embeddings.append(tf.expand_dims(values, 2))
    return maybe_dropout(concat(embeddings, axis=2), keep_prob, is_training), word_embeddings