jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def _eval_step(self, batch):
        features, input_length, phone_labels, phone_label_length, tar_label = batch
        max_length = tf.shape(tar_label)[1]
        enc_output = self.encoder(features, training=False)
        ctc_output = self.ctc_model(enc_output, training=False)

        ctc_output = tf.nn.softmax(ctc_output, -1)
        ctc_loss = tf.keras.backend.ctc_batch_cost(tf.cast(phone_labels, tf.int32),
                                                   tf.cast(ctc_output, tf.float32),
                                                   tf.cast(input_length[:, tf.newaxis], tf.int32),
                                                   tf.cast(phone_label_length[:, tf.newaxis], tf.int32),
                                                   )
        ctc_decode = tf.keras.backend.ctc_decode(ctc_output, input_length)[0][0]
        ctc_decode = tf.cast(tf.clip_by_value(ctc_decode, 0, self.phone_featurizer.num_classes), tf.int32)

        ctc_out = self.translator([ctc_decode, enc_output], training=False)


        translate_loss = tf.keras.losses.sparse_categorical_crossentropy(tar_label, ctc_out[:, :max_length], True)
        ctc_acc = self.ctc_acc(phone_labels, ctc_decode)
        translate_acc = self.translate_acc(tar_label, ctc_out, max_length)
        self.eval_metrics["ctc_loss"].update_state(ctc_loss)
        self.eval_metrics["translate_loss"].update_state(translate_loss)
        self.eval_metrics['ctc_acc'].update_state(ctc_acc)
        self.eval_metrics['translate_acc'].update_state(translate_acc)
//  [python]
def get_ragged_tensor_from_hlo(
  hlo_graph
) -> Tuple[tf.RaggedTensor, tf.RaggedTensor]:
  """Get the operands and users of each node from the hlo graph.

  Args:
    hlo_graph: a hlo graph

  Returns:
    a tuple of tf.RaggedTensor indicating the graph structure:
    (operands, users)
    operands: a dict in the form of raggedtensor, where the key is the node id,
      and the value is a list of operand indices.
    users: a dict in the form of raggedtensor, where the key is the node id,
      and the value is a list of user indices.
  """

  in_edge_features = hlo_graph.in_edge_features
  out_edge_features = hlo_graph.out_edge_features

  operands = tf.to_int64(
    tf.RaggedTensor.from_row_splits(
      values=in_edge_features.srcs, row_splits=hlo_graph.in_edge_offsets
    )
  )

  users = tf.to_int64(
    tf.RaggedTensor.from_row_splits(
      values=out_edge_features.dsts, row_splits=hlo_graph.out_edge_offsets
    )
  )

  return operands, users
//  [python]
def get_ragged_tensor_from_hlo(
  hlo_graph
) -> Tuple[tf.RaggedTensor, tf.RaggedTensor]:
  """Get the operands and users of each node from the hlo graph.

  Args:
    hlo_graph: a hlo graph

  Returns:
    a tuple of tf.RaggedTensor indicating the graph structure:
    (operands, users)
    operands: a dict in the form of raggedtensor, where the key is the node id,
      and the value is a list of operand indices.
    users: a dict in the form of raggedtensor, where the key is the node id,
      and the value is a list of user indices.
  """

  in_edge_features = hlo_graph.in_edge_features
  out_edge_features = hlo_graph.out_edge_features

  operands = tf.cast(
    tf.RaggedTensor.from_row_splits(
      values=in_edge_features.srcs, row_splits=hlo_graph.in_edge_offsets
    ), tf.int64
  )

  users = tf.cast(
    tf.RaggedTensor.from_row_splits(
      values=out_edge_features.dsts, row_splits=hlo_graph.out_edge_offsets
    ), tf.int64
  )

  return operands, users
//  [python]
def softmax_pred_op(self):
        if not self.CRF:
            self.labels_softmax_ = tf.argmax(self.logits, axis=-1)
            self.labels_softmax_ = tf.to_int32(self.labels_softmax_)
//  [python]
def softmax_pred_op(self):
        if not self.CRF:
            self.labels_softmax_ = tf.argmax(self.logits, axis=-1)
            self.labels_softmax_ = tf.cast(self.labels_softmax_, tf.int32)
//  [python]
def __init__(self,
               input_width,
               num_predictions,
               source_network,
               float_type,
               activation=None,
               initializer='glorot_uniform',
               output='logits',
               **kwargs):

    embedding_table = source_network.get_embedding_table()
    vocab_size, hidden_size = embedding_table.shape

    sequence_data = tf.keras.layers.Input(
        shape=(None, input_width), name='sequence_data', dtype=tf.float32)
    masked_lm_positions = tf.keras.layers.Input(
        shape=(num_predictions,), name='masked_lm_positions', dtype=tf.int32)

    masked_lm_input = tf.keras.layers.Lambda(
        lambda x: self._gather_indexes(x[0], x[1]))(
            [sequence_data, masked_lm_positions])
    lm_data = (
        tf.keras.layers.Dense(
            hidden_size,
            activation=activation,
            kernel_initializer=initializer,
            name='cls/predictions/transform/dense')(masked_lm_input))
    lm_data = tf.keras.layers.LayerNormalization(
        axis=-1, epsilon=1e-12, name='cls/predictions/transform/LayerNorm')(
            lm_data)
    lm_data = tf.keras.layers.Lambda(
        lambda x: tf.matmul(x, tf.to_float(embedding_table), transpose_b=True))(
            lm_data)
    logits = Bias(
        initializer=tf.keras.initializers.Zeros(),
        name='cls/predictions/output_bias')(
            lm_data)

    # We can't use the standard Keras reshape layer here, since it expects
    # the input and output batch size to be the same.
    reshape_layer = tf.keras.layers.Lambda(
        lambda x: tf.reshape(x, [-1, num_predictions, vocab_size]))

    self.logits = reshape_layer(logits)
    predictions = tf.keras.layers.Activation(tf.nn.log_softmax, dtype='float32')(self.logits)

    if output == 'logits':
      output_tensors = self.logits
    elif output == 'predictions':
      output_tensors = predictions
    else:
      raise ValueError(
          ('Unknown `output` value "%s". `output` can be either "logits" or '
           '"predictions"') % output)

    super(MaskedLM, self).__init__(
        inputs=[sequence_data, masked_lm_positions],
        outputs=output_tensors,
        **kwargs)
//  [python]
def __init__(self,
               input_width,
               num_predictions,
               source_network,
               float_type,
               activation=None,
               initializer='glorot_uniform',
               output='logits',
               **kwargs):

    embedding_table = source_network.get_embedding_table()
    vocab_size, hidden_size = embedding_table.shape

    sequence_data = tf.keras.layers.Input(
        shape=(None, input_width), name='sequence_data', dtype=tf.float32)
    masked_lm_positions = tf.keras.layers.Input(
        shape=(num_predictions,), name='masked_lm_positions', dtype=tf.int32)

    masked_lm_input = tf.keras.layers.Lambda(
        lambda x: self._gather_indexes(x[0], x[1]))(
            [sequence_data, masked_lm_positions])
    lm_data = (
        tf.keras.layers.Dense(
            hidden_size,
            activation=activation,
            kernel_initializer=initializer,
            name='cls/predictions/transform/dense')(masked_lm_input))
    lm_data = tf.keras.layers.LayerNormalization(
        axis=-1, epsilon=1e-12, name='cls/predictions/transform/LayerNorm')(
            lm_data)
    lm_data = tf.keras.layers.Lambda(
        lambda x: tf.matmul(x, tf.cast(embedding_table, float_type), transpose_b=True))(
            lm_data)
    logits = Bias(
        initializer=tf.keras.initializers.Zeros(),
        name='cls/predictions/output_bias')(
            lm_data)

    # We can't use the standard Keras reshape layer here, since it expects
    # the input and output batch size to be the same.
    reshape_layer = tf.keras.layers.Lambda(
        lambda x: tf.reshape(x, [-1, num_predictions, vocab_size]))

    self.logits = reshape_layer(logits)
    predictions = tf.keras.layers.Activation(tf.nn.log_softmax, dtype='float32')(self.logits)

    if output == 'logits':
      output_tensors = self.logits
    elif output == 'predictions':
      output_tensors = predictions
    else:
      raise ValueError(
          ('Unknown `output` value "%s". `output` can be either "logits" or '
           '"predictions"') % output)

    super(MaskedLM, self).__init__(
        inputs=[sequence_data, masked_lm_positions],
        outputs=output_tensors,
        **kwargs)
//  [python]
def get_error(prediction,
              label,
              label_length,
              out_charset,
              name,
              max_text_summary=10):
    """ Returns the wrong number between prediction and label.
    """
    # Calculate Edit dist.
    hypothesis = t.cast(prediction, t.int32)
    edit_distance = tf.edit_distance(hypothesis, label, normalize=False)

    # Count wrong seq.
    error_count = tf.count_nonzero(edit_distance, axis=0)

    # Summary predicted text
    num_classes = len(out_charset)
    prediction = tf.sparse_tensor_to_dense(prediction,
                                           default_value=num_classes)
    label = tf.sparse_tensor_to_dense(label, default_value=num_classes)
    table = get_itos_table(out_charset)
    pr_text = tf.reduce_join(table.lookup(tf.to_int64(prediction)),
                             1,
                             keep_dims=True)
    la_text = tf.reduce_join(table.lookup(tf.to_int64(label)),
                             1,
                             keep_dims=True)
    tf.summary.text('label', la_text[:max_text_summary])
    tf.summary.text('%s/predicted' % name, pr_text[:max_text_summary])

    return error_count
//  [python]
def get_error(prediction,
              label,
              label_length,
              out_charset,
              name,
              max_text_summary=10):
    """ Returns the wrong number between prediction and label.
    """
    # Calculate Edit dist.
    hypothesis = tf.cast(prediction, tf.int32)
    edit_distance = tf.edit_distance(hypothesis, label, normalize=False)

    # Count wrong seq.
    error_count = tf.count_nonzero(edit_distance, axis=0)

    # Summary predicted text
    num_classes = len(out_charset)
    prediction = tf.sparse_tensor_to_dense(prediction,
                                           default_value=num_classes)
    label = tf.sparse_tensor_to_dense(label, default_value=num_classes)
    table = get_itos_table(out_charset)
    pr_text = tf.reduce_join(table.lookup(tf.to_int64(prediction)),
                             1,
                             keep_dims=True)
    la_text = tf.reduce_join(table.lookup(tf.to_int64(label)),
                             1,
                             keep_dims=True)
    tf.summary.text('label', la_text[:max_text_summary])
    tf.summary.text('%s/predicted' % name, pr_text[:max_text_summary])

    return error_count
//  [python]
def _train_step(self, batch):
        features, input_length, phone_labels, phone_label_length, tar_label = batch

        max_length=tf.shape(tar_label)[1]
        with tf.GradientTape() as tape:

            enc_output = self.encoder(features, training=True)
            ctc_output = self.ctc_model(enc_output, training=True)
            #print(enc_output.shape,ctc_output.shape,phone_labels.shape,input_length.shape,phone_label_length.shape)

            ctc_output = tf.nn.softmax(ctc_output, -1)
            ctc_loss = tf.keras.backend.ctc_batch_cost(t.cast(phone_labels, tf.int32),
                                            t.cast(ctc_output, tf.float32),
                                            t.cast(input_length[:,tf.newaxis], tf.int32),
                                            t.cast(phone_label_length[:,tf.newaxis], tf.int32),
                                            )

            ctc_decode_result=tf.keras.backend.ctc_decode(ctc_output,input_length)[0][0]
            ctc_decode_result=t.cast(tf.clip_by_value(ctc_decode_result,0,self.phone_featurizer.num_classes),tf.int32)
            label_out = self.translator([t.concat([phone_labels,t.zeros_like(phone_labels)[:,:5]],-1),enc_output], training=True)
            ctc_out = self.translator([ctc_decode_result,enc_output], training=True)

            translate_loss=self.mask_loss(tar_label,label_out[:,:max_length])*2.+ self.mask_loss(tar_label, ctc_out[:,:max_length])
            # train_loss = tf.reduce_mean(ctc_loss+translate_loss*5.)
            train_loss=tf.nn.compute_average_loss(ctc_loss+translate_loss*2.,global_batch_size=self.global_batch_size)

            if self.is_mixed_precision:
                scaled_train_loss = self.optimizer.get_scaled_loss(train_loss)
        train_list=self.encoder.trainable_variables+self.ctc_model.trainable_variables+self.translator.trainable_variables
        if self.is_mixed_precision:
            scaled_gradients = tape.gradient(scaled_train_loss, train_list)
            gradients = self.optimizer.get_unscaled_gradients(scaled_gradients)
        else:
            gradients = tape.gradient(train_loss, train_list)
        self.optimizer.apply_gradients(zip(gradients, train_list))

        ctc_acc = self.ctc_acc(phone_labels, ctc_decode_result)
        translate_acc=self.translate_acc(tar_label,ctc_out,max_length)
        self.train_metrics['ctc_acc'].update_state(ctc_acc)
        self.train_metrics["ctc_loss"].update_state(ctc_loss)
        self.train_metrics["train_loss"].update_state(train_loss)
        self.train_metrics["translate_loss"].update_state(translate_loss)
        self.train_metrics["translate_acc"].update_state(translate_acc)
//  [python]
def _train_step(self, batch):
        features, input_length, phone_labels, phone_label_length, tar_label = batch

        max_length=tf.shape(tar_label)[1]
        with tf.GradientTape() as tape:

            enc_output = self.encoder(features, training=True)
            ctc_output = self.ctc_model(enc_output, training=True)
            #print(enc_output.shape,ctc_output.shape,phone_labels.shape,input_length.shape,phone_label_length.shape)

            ctc_output = tf.nn.softmax(ctc_output, -1)
            ctc_loss = tf.keras.backend.ctc_batch_cost(tf.cast(phone_labels, tf.int32),
                                            tf.cast(ctc_output, tf.float32),
                                            tf.cast(input_length[:,tf.newaxis], tf.int32),
                                            tf.cast(phone_label_length[:,tf.newaxis], tf.int32),
                                            )

            ctc_decode_result=tf.keras.backend.ctc_decode(ctc_output,input_length)[0][0]
            ctc_decode_result=tf.cast(tf.clip_by_value(ctc_decode_result,0,self.phone_featurizer.num_classes),tf.int32)
            label_out = self.translator([tf.concat([phone_labels,tf.zeros_like(phone_labels)[:,:5]],-1),enc_output], training=True)
            ctc_out = self.translator([ctc_decode_result,enc_output], training=True)

            translate_loss=self.mask_loss(tar_label,label_out[:,:max_length])*2.+ self.mask_loss(tar_label, ctc_out[:,:max_length])
            # train_loss = tf.reduce_mean(ctc_loss+translate_loss*5.)
            train_loss=tf.nn.compute_average_loss(ctc_loss+translate_loss*2.,global_batch_size=self.global_batch_size)

            if self.is_mixed_precision:
                scaled_train_loss = self.optimizer.get_scaled_loss(train_loss)
        train_list=self.encoder.trainable_variables+self.ctc_model.trainable_variables+self.translator.trainable_variables
        if self.is_mixed_precision:
            scaled_gradients = tape.gradient(scaled_train_loss, train_list)
            gradients = self.optimizer.get_unscaled_gradients(scaled_gradients)
        else:
            gradients = tape.gradient(train_loss, train_list)
        self.optimizer.apply_gradients(zip(gradients, train_list))

        ctc_acc = self.ctc_acc(phone_labels, ctc_decode_result)
        translate_acc=self.translate_acc(tar_label,ctc_out,max_length)
        self.train_metrics['ctc_acc'].update_state(ctc_acc)
        self.train_metrics["ctc_loss"].update_state(ctc_loss)
        self.train_metrics["train_loss"].update_state(train_loss)
        self.train_metrics["translate_loss"].update_state(translate_loss)
        self.train_metrics["translate_acc"].update_state(translate_acc)
//  [python]
def central_crop(image, crop_size):
  """Returns a central crop for the specified size of an image.

  Args:
    image: A tensor with shape [height, width, channels]
    crop_size: A tuple (crop_width, crop_height)

  Returns:
    A tensor of shape [crop_height, crop_width, channels].
  """
  with tf.compat.v1.variable_scope('CentralCrop'):
    target_width, target_height = crop_size
    image_height, image_width = tf.shape(
        input=image)[0], tf.shape(input=image)[1]
    assert_op1 = tf.Assert(
        tf.greater_equal(image_height, target_height),
        ['image_height < target_height', image_height, target_height])
    assert_op2 = tf.Assert(
        tf.greater_equal(image_width, target_width),
        ['image_width < target_width', image_width, target_width])
    with tf.control_dependencies([assert_op1, assert_op2]):
      offset_width = tf.to_int32((image_width - target_width) / 2)
      offset_height = tf.to_int32((image_height - target_height) / 2)
      return tf.image.crop_to_bounding_box(image, offset_height, offset_width,
                                           target_height, target_width)
//  [python]
def central_crop(image, crop_size):
  """Returns a central crop for the specified size of an image.

  Args:
    image: A tensor with shape [height, width, channels]
    crop_size: A tuple (crop_width, crop_height)

  Returns:
    A tensor of shape [crop_height, crop_width, channels].
  """
  with tf.compat.v1.variable_scope('CentralCrop'):
    target_width, target_height = crop_size
    image_height, image_width = tf.shape(
        input=image)[0], tf.shape(input=image)[1]
    assert_op1 = tf.Assert(
        tf.greater_equal(image_height, target_height),
        ['image_height < target_height', image_height, target_height])
    assert_op2 = tf.Assert(
        tf.greater_equal(image_width, target_width),
        ['image_width < target_width', image_width, target_width])
    with tf.control_dependencies([assert_op1, assert_op2]):
      offset_width = tf.cast((image_width - target_width) / 2, tf.int32)
      offset_height = tf.cast((image_height - target_height) / 2, tf.int32)
      return tf.image.crop_to_bounding_box(image, offset_height, offset_width,
                                           target_height, target_width)
//  [python]
def accur_1(y_true, y_pred):
    y_true = tf.round(y_true)
    pred = tf.round(y_pred)
    summ = tf.constant(MAX_BITMAP_SIZE, dtype=tf.float32)
    wrong_num = tf.subtract(summ, tf.reduce_sum(tf.to_float(tf.equal(y_true, pred)), axis=-1))
    right_1_num = tf.reduce_sum(tf.to_float(tf.logical_and(tf.to_bool(y_true), tf.to_bool(pred))), axis=-1)
    return K.mean(tf.divide(right_1_num, tf.add(right_1_num, wrong_num)))
//  [python]
def accur_1(y_true, y_pred):
    y_true = tf.round(y_true)
    pred = tf.round(y_pred)
    summ = tf.constant(MAX_BITMAP_SIZE, dtype=tf.float32)
    wrong_num = tf.subtract(summ, tf.reduce_sum(tf.cast(tf.equal(y_true, pred), tf.float32), axis=-1))
    right_1_num = tf.reduce_sum(tf.cast(tf.logical_and(tf.cast(y_true, tf.bool), tf.cast(pred, tf.bool)), tf.float32), axis=-1)
    return K.mean(tf.divide(right_1_num, tf.add(right_1_num, wrong_num)))
//  [python]
def mean(x, axis=None, keepdims=False):
    """Mean of a tensor, alongside the specified axis.

    # Arguments
        x: A tensor or variable.
        axis: A list of integer. Axes to compute the mean.
        keepdims: A boolean, whether to keep the dimensions or not.
            If `keepdims` is `False`, the rank of the tensor is reduced
            by 1 for each entry in `axis`. If `keep_dims` is `True`,
            the reduced dimensions are retained with length 1.

    # Returns
        A tensor with the mean of elements of `x`.
    """
    axis = _normalize_axis(axis, ndim(x))
    if x.dtype.base_dtype == tf.bool:
        x = t.cast(x, floatx())
    return tf.reduce_mean(x, axis=axis, keep_dims=keepdims)
//  [python]
def mean(x, axis=None, keepdims=False):
    """Mean of a tensor, alongside the specified axis.

    # Arguments
        x: A tensor or variable.
        axis: A list of integer. Axes to compute the mean.
        keepdims: A boolean, whether to keep the dimensions or not.
            If `keepdims` is `False`, the rank of the tensor is reduced
            by 1 for each entry in `axis`. If `keep_dims` is `True`,
            the reduced dimensions are retained with length 1.

    # Returns
        A tensor with the mean of elements of `x`.
    """
    axis = _normalize_axis(axis, ndim(x))
    if x.dtype.base_dtype == tf.bool:
        x = tf.cast(x, floatx())
    return tf.reduce_mean(x, axis=axis, keep_dims=keepdims)
//  [python]
def var(x, axis=None, keepdims=False):
    """Variance of a tensor, alongside the specified axis.

    # Arguments
        x: A tensor or variable.
        axis: An integer, the axis to compute the variance.
        keepdims: A boolean, whether to keep the dimensions or not.
            If `keepdims` is `False`, the rank of the tensor is reduced
            by 1. If `keepdims` is `True`,
            the reduced dimension is retained with length 1.

    # Returns
        A tensor with the variance of elements of `x`.
    """
    axis = _normalize_axis(axis, ndim(x))
    if x.dtype.base_dtype == tf.bool:
        x = t.cast(x, floatx())
    m = tf.reduce_mean(x, axis=axis, keep_dims=True)
    devs_squared = tf.square(x - m)
    return tf.reduce_mean(devs_squared,
                          axis=axis,
                          keepdims=keepdims)
//  [python]
def var(x, axis=None, keepdims=False):
    """Variance of a tensor, alongside the specified axis.

    # Arguments
        x: A tensor or variable.
        axis: An integer, the axis to compute the variance.
        keepdims: A boolean, whether to keep the dimensions or not.
            If `keepdims` is `False`, the rank of the tensor is reduced
            by 1. If `keepdims` is `True`,
            the reduced dimension is retained with length 1.

    # Returns
        A tensor with the variance of elements of `x`.
    """
    axis = _normalize_axis(axis, ndim(x))
    if x.dtype.base_dtype == tf.bool:
        x = tf.cast(x, floatx())
    m = tf.reduce_mean(x, axis=axis, keep_dims=True)
    devs_squared = tf.square(x - m)
    return tf.reduce_mean(devs_squared,
                          axis=axis,
                          keep_dims=keepdims)
//  [python]
def any(x, axis=None, keepdims=False):
    """Bitwise reduction (logical OR).

    # Arguments
        x: Tensor or variable.
        axis: axis along which to perform the reduction.
        keepdims: whether the drop or broadcast the reduction axes.

    # Returns
        A uint8 tensor (0s and 1s).
    """
    axis = _normalize_axis(axis, ndim(x))
    x = t.cast(x, t.bool)
    return tf.reduce_any(x, axis=axis, keep_dims=keepdims)
//  [python]
def any(x, axis=None, keepdims=False):
    """Bitwise reduction (logical OR).

    # Arguments
        x: Tensor or variable.
        axis: axis along which to perform the reduction.
        keepdims: whether the drop or broadcast the reduction axes.

    # Returns
        A uint8 tensor (0s and 1s).
    """
    axis = _normalize_axis(axis, ndim(x))
    x = tf.cast(x, tf.bool)
    return tf.reduce_any(x, axis=axis, keep_dims=keepdims)
//  [python]
def all(x, axis=None, keepdims=False):
    """Bitwise reduction (logical AND).

    # Arguments
        x: Tensor or variable.
        axis: axis along which to perform the reduction.
        keepdims: whether the drop or broadcast the reduction axes.

    # Returns
        A uint8 tensor (0s and 1s).
    """
    axis = _normalize_axis(axis, ndim(x))
    x = t.cast(x, t.bool)
    return tf.reduce_all(x, axis=axis, keep_dims=keepdims)
//  [python]
def all(x, axis=None, keepdims=False):
    """Bitwise reduction (logical AND).

    # Arguments
        x: Tensor or variable.
        axis: axis along which to perform the reduction.
        keepdims: whether the drop or broadcast the reduction axes.

    # Returns
        A uint8 tensor (0s and 1s).
    """
    axis = _normalize_axis(axis, ndim(x))
    x = tf.cast(x, tf.bool)
    return tf.reduce_all(x, axis=axis, keep_dims=keepdims)
//  [python]
def call(self, inputs: tf.Tensor, mask: tf.Tensor):

        rmask = tf.logical_not(t.cast(mask, tf.bool))
        output = tf.where(rmask, float("-inf"), inputs)
        output = stable_softmax(output, self.axis)
        output = tf.where(rmask, 0.0, output)
        return output
//  [python]
def call(self, inputs: tf.Tensor, mask: tf.Tensor):

        rmask = tf.logical_not(tf.cast(mask, tf.bool))
        output = tf.where(rmask, float("-inf"), inputs)
        output = stable_softmax(output, self.axis)
        output = tf.where(rmask, 0.0, output)
        return output
//  [python]
def cmpm_loss_compute(text_embeddings, image_embeddings, labels):
    """ Cross-Modal Projection Matching Loss (CMPM)

    Args:
        text_embeddings: text joint embeddings
        image_embeddings: image joint embeddings
        labels: class labels

    Returns:
        i2t_matching_loss: cmpm loss for image projected to text
        t2i_matching_loss: cmpm loss for text projected to image
        pos_avg_dist: average distance of positive pairs
        neg_avg_dist: average distance of negative pairs
    """
    # label mask
    batch_size = image_embeddings.get_shape().as_list()[0]
    mylabels = tf.to_float(tf.reshape(labels, [batch_size, 1]))
    labelD = pairwise_distance(mylabels, mylabels)
    label_mask = tf.to_float(tf.less(labelD, 0.5))  # 1-match   0-unmatch

    # cross-modal scalar projection
    image_embeddings_norm = tf.nn.l2_normalize(image_embeddings, dim=-1)
    text_embeddings_norm = tf.nn.l2_normalize(text_embeddings, dim=-1)

    image_proj_text = tf.matmul(image_embeddings, tf.transpose(text_embeddings_norm))
    text_proj_image = tf.matmul(text_embeddings, tf.transpose(image_embeddings_norm))

    # softmax, higher scalar projection gives higher probability
    i2t_pred = tf.nn.softmax(image_proj_text)
    t2i_pred = tf.nn.softmax(text_proj_image)

    # normalize the true matching distribution
    label_mask = tf.divide(label_mask, tf.reduce_sum(label_mask, axis=1, keep_dims=True))

    # KL Divergence
    i2t_matching_loss = tf.reduce_mean(tf.reduce_sum(i2t_pred * tf.log(1e-8 + i2t_pred / (label_mask + 1e-8)), 1))
    t2i_matching_loss = tf.reduce_mean(tf.reduce_sum(t2i_pred * tf.log(1e-8 + t2i_pred / (label_mask + 1e-8)), 1))

    # averaged cosine distance of positive and negative pairs for observation
    cosdist = 1.0 - tf.matmul(text_embeddings_norm, tf.transpose(image_embeddings_norm))

    pos_avg_dist = tf.reduce_mean(tf.boolean_mask(cosdist, tf.less(labelD, 0.5)))
    neg_avg_dist = tf.reduce_mean(tf.boolean_mask(cosdist, tf.greater(labelD, 0.5)))

    return i2t_matching_loss, t2i_matching_loss, pos_avg_dist, neg_avg_dist
//  [python]
def cmpm_loss_compute(text_embeddings, image_embeddings, labels):
    """ Cross-Modal Projection Matching Loss (CMPM)

    Args:
        text_embeddings: text joint embeddings
        image_embeddings: image joint embeddings
        labels: class labels

    Returns:
        i2t_matching_loss: cmpm loss for image projected to text
        t2i_matching_loss: cmpm loss for text projected to image
        pos_avg_dist: average distance of positive pairs
        neg_avg_dist: average distance of negative pairs
    """
    # label mask
    batch_size = image_embeddings.get_shape().as_list()[0]
    mylabels = tf.cast(tf.reshape(labels, [batch_size, 1]), tf.float32)
    labelD = pairwise_distance(mylabels, mylabels)
    label_mask = tf.cast(tf.less(labelD, 0.5), tf.float32)  # 1-match   0-unmatch

    # cross-modal scalar projection
    image_embeddings_norm = tf.nn.l2_normalize(image_embeddings, dim=-1)
    text_embeddings_norm = tf.nn.l2_normalize(text_embeddings, dim=-1)

    image_proj_text = tf.matmul(image_embeddings, tf.transpose(text_embeddings_norm))
    text_proj_image = tf.matmul(text_embeddings, tf.transpose(image_embeddings_norm))

    # softmax, higher scalar projection gives higher probability
    i2t_pred = tf.nn.softmax(image_proj_text)
    t2i_pred = tf.nn.softmax(text_proj_image)

    # normalize the true matching distribution
    label_mask = tf.divide(label_mask, tf.reduce_sum(label_mask, axis=1, keep_dims=True))

    # KL Divergence
    i2t_matching_loss = tf.reduce_mean(tf.reduce_sum(i2t_pred * tf.log(1e-8 + i2t_pred / (label_mask + 1e-8)), 1))
    t2i_matching_loss = tf.reduce_mean(tf.reduce_sum(t2i_pred * tf.log(1e-8 + t2i_pred / (label_mask + 1e-8)), 1))

    # averaged cosine distance of positive and negative pairs for observation
    cosdist = 1.0 - tf.matmul(text_embeddings_norm, tf.transpose(image_embeddings_norm))

    pos_avg_dist = tf.reduce_mean(tf.boolean_mask(cosdist, tf.less(labelD, 0.5)))
    neg_avg_dist = tf.reduce_mean(tf.boolean_mask(cosdist, tf.greater(labelD, 0.5)))

    return i2t_matching_loss, t2i_matching_loss, pos_avg_dist, neg_avg_dist
//  [python]
def xdropout(self, inputs):
        """
        Applies dropout to the inputs, as vanilla dropout, but also scales the remaining elements up by 1/drop_prob.
        """
        mask = tf.cast(
            1
            - tf.compat.v1.distributions.Bernoulli(probs=1.0 - self.drop_prob).sample(sample_shape=shape_list(inputs)),
            tf.bool,
        )
        scale = tf.convert_to_tensor(1.0 / (1 - self.drop_prob), dtype=tf.float32)
        if self.drop_prob > 0:
            inputs = tf.where(mask, 0.0, inputs) * scale

        def grad(upstream):
            if self.drop_prob > 0:
                return tf.where(mask, 0.0, upstream) * scale
            else:
                return upstream

        return inputs, grad
//  [python]
def get_attention_mask(self, attention_mask):
        if len(shape_list(attention_mask)) <= 2:
            extended_attention_mask = tf.expand_dims(tf.expand_dims(attention_mask, 1), 2)
            attention_mask = extended_attention_mask * tf.expand_dims(tf.squeeze(extended_attention_mask, -2), -1)
            attention_mask = tf.to_int32(attention_mask)
        elif len(shape_list(attention_mask)) == 3:
            attention_mask = tf.expand_dims(attention_mask, 1)

        return attention_mask
//  [python]
def get_attention_mask(self, attention_mask):
        if len(shape_list(attention_mask)) <= 2:
            extended_attention_mask = tf.expand_dims(tf.expand_dims(attention_mask, 1), 2)
            attention_mask = extended_attention_mask * tf.expand_dims(tf.squeeze(extended_attention_mask, -2), -1)
            attention_mask = tf.cast(attention_mask, tf.uint8)
        elif len(shape_list(attention_mask)) == 3:
            attention_mask = tf.expand_dims(attention_mask, 1)

        return attention_mask
//  [python]
def ctc_label_dense_to_sparse(labels, label_lengths):
    """Converts CTC labels from dense to sparse.

    # Arguments
        labels: dense CTC labels.
        label_lengths: length of the labels.

    # Returns
        A sparse tensor representation of the lablels.
    """
    label_shape = tf.shape(labels)
    num_batches_tns = tf.stack([label_shape[0]])
    max_num_labels_tns = tf.stack([label_shape[1]])

    def range_less_than(_, current_input):
        return tf.expand_dims(tf.range(label_shape[1]), 0) < tf.fill(
            max_num_labels_tns, current_input)

    init = tf.fill([1, label_shape[1]], 0)
    dense_mask = functional_ops.scan(range_less_than, label_lengths,
                                     initializer=init, parallel_iterations=1)
    dense_mask = dense_mask[:, 0, :]

    label_array = tf.reshape(tf.tile(tf.range(0, label_shape[1]), num_batches_tns),
                             label_shape)
    label_ind = tf.boolean_mask(label_array, dense_mask)

    batch_array = tf.transpose(tf.reshape(tf.tile(tf.range(0, label_shape[0]),
                                                  max_num_labels_tns), reverse(label_shape, 0)))
    batch_ind = tf.boolean_mask(batch_array, dense_mask)
    indices = tf.transpose(tf.reshape(concatenate([batch_ind, label_ind], axis=0), [2, -1]))

    vals_sparse = tf.gather_nd(labels, indices)

    return tf.SparseTensor(tf.to_int64(indices), vals_sparse, tf.to_int64(label_shape))
//  [python]
def ctc_label_dense_to_sparse(labels, label_lengths):
    """Converts CTC labels from dense to sparse.

    # Arguments
        labels: dense CTC labels.
        label_lengths: length of the labels.

    # Returns
        A sparse tensor representation of the lablels.
    """
    label_shape = tf.shape(labels)
    num_batches_tns = tf.stack([label_shape[0]])
    max_num_labels_tns = tf.stack([label_shape[1]])

    def range_less_than(_, current_input):
        return tf.expand_dims(tf.range(label_shape[1]), 0) < tf.fill(
            max_num_labels_tns, current_input)

    init = tf.cast(tf.fill([1, label_shape[1]], 0), tf.bool)
    dense_mask = functional_ops.scan(range_less_than, label_lengths,
                                     initializer=init, parallel_iterations=1)
    dense_mask = dense_mask[:, 0, :]

    label_array = tf.reshape(tf.tile(tf.range(0, label_shape[1]), num_batches_tns),
                             label_shape)
    label_ind = tf.boolean_mask(label_array, dense_mask)

    batch_array = tf.transpose(tf.reshape(tf.tile(tf.range(0, label_shape[0]),
                                                  max_num_labels_tns), reverse(label_shape, 0)))
    batch_ind = tf.boolean_mask(batch_array, dense_mask)
    indices = tf.transpose(tf.reshape(concatenate([batch_ind, label_ind], axis=0), [2, -1]))

    vals_sparse = tf.gather_nd(labels, indices)

    return tf.SparseTensor(tf.to_int64(indices), vals_sparse, tf.to_int64(label_shape))
//  [python]
def call(
        self, hidden_states: tf.Tensor, residual_states: tf.Tensor, input_mask: tf.Tensor, training: bool = False
    ) -> tf.Tensor:
        out = tf.nn.conv2d(
            tf.expand_dims(hidden_states, 1),
            tf.expand_dims(self.conv_kernel, 0),
            strides=1,
            padding=[[0, 0], [0, 0], [self.padding, self.padding], [0, 0]],
        )
        out = tf.squeeze(tf.nn.bias_add(out, self.conv_bias), 1)
        rmask = tf.to_float(1 - input_mask) > 0
        out = tf.where(tf.broadcast_to(tf.expand_dims(rmask, -1), shape_list(out)), 0.0, out)
        out = self.dropout(out, training=training)
        out = self.conv_act(out)

        layer_norm_input = residual_states + out
        output = self.LayerNorm(layer_norm_input)

        if input_mask is None:
            output_states = output
        else:
            if len(shape_list(input_mask)) != len(shape_list(layer_norm_input)):
                if len(shape_list(input_mask)) == 4:
                    input_mask = tf.squeeze(tf.squeeze(input_mask, axis=1), axis=1)
                input_mask = tf.to_float(tf.expand_dims(input_mask, axis=2))

            output_states = output * input_mask

        return output_states
//  [python]
def call(
        self, hidden_states: tf.Tensor, residual_states: tf.Tensor, input_mask: tf.Tensor, training: bool = False
    ) -> tf.Tensor:
        out = tf.nn.conv2d(
            tf.expand_dims(hidden_states, 1),
            tf.expand_dims(self.conv_kernel, 0),
            strides=1,
            padding=[[0, 0], [0, 0], [self.padding, self.padding], [0, 0]],
        )
        out = tf.squeeze(tf.nn.bias_add(out, self.conv_bias), 1)
        rmask = tf.cast(1 - input_mask, tf.bool)
        out = tf.where(tf.broadcast_to(tf.expand_dims(rmask, -1), shape_list(out)), 0.0, out)
        out = self.dropout(out, training=training)
        out = self.conv_act(out)

        layer_norm_input = residual_states + out
        output = self.LayerNorm(layer_norm_input)

        if input_mask is None:
            output_states = output
        else:
            if len(shape_list(input_mask)) != len(shape_list(layer_norm_input)):
                if len(shape_list(input_mask)) == 4:
                    input_mask = tf.squeeze(tf.squeeze(input_mask, axis=1), axis=1)
                input_mask = tf.cast(tf.expand_dims(input_mask, axis=2), tf.float32)

            output_states = output * input_mask

        return output_states
//  [python]
def make_log_bucket_position(relative_pos, bucket_size, max_position):
    sign = tf.math.sign(relative_pos)
    mid = bucket_size // 2
    abs_pos = tf.where((relative_pos < mid) & (relative_pos > -mid), mid - 1, tf.math.abs(relative_pos))
    log_pos = (
        tf.math.ceil(
            tf.log(tf.to_float(tf.math.log(abs_pos / mid))) / tf.log(tf.to_float((max_position - 1) / mid)) * (mid - 1)
        )
        + mid
    )
    bucket_pos = tf.to_int32(
        tf.where(abs_pos <= mid, tf.to_float(relative_pos), log_pos * tf.to_float(sign))
    )
    return bucket_pos
//  [python]
def make_log_bucket_position(relative_pos, bucket_size, max_position):
    sign = tf.math.sign(relative_pos)
    mid = bucket_size // 2
    abs_pos = tf.where((relative_pos < mid) & (relative_pos > -mid), mid - 1, tf.math.abs(relative_pos))
    log_pos = (
        tf.math.ceil(
            tf.cast(tf.math.log(abs_pos / mid), tf.float32) / tf.math.log((max_position - 1) / mid) * (mid - 1)
        )
        + mid
    )
    bucket_pos = tf.cast(
        tf.where(abs_pos <= mid, tf.cast(relative_pos, tf.float32), log_pos * tf.cast(sign, tf.float32)), tf.int32
    )
    return bucket_pos
//  [python]
def assign(self, variable, value):
        if isinstance(variable, KerasVariable):
            variable = variable.value
        value = tf.cast(value, variable.dtype)
        if isinstance(value, tf.IndexedSlices):
            variable.scatter_update(value)
        else:
            variable.assign(value)
//  [python]
def assign_add(self, variable, value):
        if isinstance(variable, KerasVariable):
            variable = variable.value
        value = tf.cast(value, variable.dtype)
        if isinstance(value, tf.IndexedSlices):
            variable.scatter_add(value)
        else:
            variable.assign_add(value)
//  [python]
def assign_sub(self, variable, value):
        if isinstance(variable, KerasVariable):
            variable = variable.value
        value = tf.cast(value, variable.dtype)
        if isinstance(value, tf.IndexedSlices):
            variable.scatter_sub(value)
        else:
            variable.assign_sub(value)
//  [python]
def build_relative_position(query_size, key_size, bucket_size=-1, max_position=-1):
    """
    Build relative position according to the query and key

    We assume the absolute position of query \\(P_q\\) is range from (0, query_size) and the absolute position of key
    \\(P_k\\) is range from (0, key_size), The relative positions from query to key is \\(R_{q \\rightarrow k} = P_q -
    P_k\\)

    Args:
        query_size (int): the length of query
        key_size (int): the length of key
        bucket_size (int): the size of position bucket
        max_position (int): the maximum allowed absolute position

    Return:
        `tf.Tensor`: A tensor with shape [1, query_size, key_size]

    """
    q_ids = tf.range(query_size, dtype=tf.int32)
    k_ids = tf.range(key_size, dtype=tf.int32)
    rel_pos_ids = q_ids[:, None] - tf.tile(tf.expand_dims(k_ids, axis=0), [shape_list(q_ids)[0], 1])
    if bucket_size > 0 and max_position > 0:
        rel_pos_ids = make_log_bucket_position(rel_pos_ids, bucket_size, max_position)
    rel_pos_ids = rel_pos_ids[:query_size, :]
    rel_pos_ids = tf.expand_dims(rel_pos_ids, axis=0)
    return t.convert_to_tensor(rel_pos_ids, dtype=tf.int64)
//  [python]
def build_relative_position(query_size, key_size, bucket_size=-1, max_position=-1):
    """
    Build relative position according to the query and key

    We assume the absolute position of query \\(P_q\\) is range from (0, query_size) and the absolute position of key
    \\(P_k\\) is range from (0, key_size), The relative positions from query to key is \\(R_{q \\rightarrow k} = P_q -
    P_k\\)

    Args:
        query_size (int): the length of query
        key_size (int): the length of key
        bucket_size (int): the size of position bucket
        max_position (int): the maximum allowed absolute position

    Return:
        `tf.Tensor`: A tensor with shape [1, query_size, key_size]

    """
    q_ids = tf.range(query_size, dtype=tf.int32)
    k_ids = tf.range(key_size, dtype=tf.int32)
    rel_pos_ids = q_ids[:, None] - tf.tile(tf.expand_dims(k_ids, axis=0), [shape_list(q_ids)[0], 1])
    if bucket_size > 0 and max_position > 0:
        rel_pos_ids = make_log_bucket_position(rel_pos_ids, bucket_size, max_position)
    rel_pos_ids = rel_pos_ids[:query_size, :]
    rel_pos_ids = tf.expand_dims(rel_pos_ids, axis=0)
    return tf.cast(rel_pos_ids, tf.int64)
//  [python]
def _apply_weight_decay(self, variables):
        if self.weight_decay is None:
            return

        def distributed_apply_weight_decay(distribution, variables, **kwargs):
            def weight_decay_fn(variable):
                if self._use_weight_decay(variable):
                    lr = tf.cast(self.learning_rate, variable.dtype)
                    wd = tf.cast(self.weight_decay, variable.dtype)
                    variable.assign_sub(variable * wd * lr)

            for variable in variables:
                if isinstance(variable, backend.Variable):
                    variable = variable.value  # Convert to tf.Variable
                distribution.extended.update(
                    variable, weight_decay_fn, group=False
                )

        tf.__internal__.distribute.interim.maybe_merge_call(
            distributed_apply_weight_decay,
            self._distribution_strategy,
            variables,
        )
//  [python]
def call(
        self,
        hidden_states: tf.Tensor,
        attention_mask: tf.Tensor,
        query_states: tf.Tensor = None,
        relative_pos: tf.Tensor = None,
        output_attentions: bool = False,
        output_hidden_states: bool = False,
        return_dict: bool = True,
        training: bool = False,
    ) -> Union[TFBaseModelOutput, Tuple[tf.Tensor]]:
        if len(shape_list(attention_mask)) <= 2:
            input_mask = attention_mask
        else:
            input_mask = tf.cast(tf.math.reduce_sum(attention_mask, axis=-2) > 0, dtype=tf.uint8)

        all_hidden_states = () if output_hidden_states else None
        all_attentions = () if output_attentions else None

        attention_mask = self.get_attention_mask(attention_mask)
        relative_pos = self.get_rel_pos(hidden_states, query_states, relative_pos)

        next_kv = hidden_states

        rel_embeddings = self.get_rel_embedding()
        output_states = next_kv
        for i, layer_module in enumerate(self.layer):

            if output_hidden_states:
                all_hidden_states = all_hidden_states + (output_states,)

            layer_outputs = layer_module(
                hidden_states=next_kv,
                attention_mask=attention_mask,
                query_states=query_states,
                relative_pos=relative_pos,
                rel_embeddings=rel_embeddings,
                output_attentions=output_attentions,
                training=training,
            )
            output_states = layer_outputs[0]

            if i == 0 and self.conv is not None:
                output_states = self.conv(hidden_states, output_states, input_mask)

            next_kv = output_states

            if output_attentions:
                all_attentions = all_attentions + (layer_outputs[1],)

        # Add last layer
        if output_hidden_states:
            all_hidden_states = all_hidden_states + (output_states,)

        if not return_dict:
            return tuple(v for v in [output_states, all_hidden_states, all_attentions] if v is not None)

        return TFBaseModelOutput(
            last_hidden_state=output_states, hidden_states=all_hidden_states, attentions=all_attentions
        )
//  [python]
def _enclosing_power_of_2(input_len: Union[int, tf.Tensor]) -> tf.Tensor:
  """Returns the smallest power of 2 greater than or equal to input_len."""
  n = tf.math.ceil(_log2(t.cast(input_len, tf.float32) - 0.5))
  return t.cast(tf.round(tf.pow(2., t.cast(n, tf.float32))), tf.int32)
//  [python]
def _enclosing_power_of_2(input_len: Union[int, tf.Tensor]) -> tf.Tensor:
  """Returns the smallest power of 2 greater than or equal to input_len."""
  n = tf.math.ceil(_log2(tf.cast(input_len, tf.float32) - 0.5))
  return tf.cast(tf.round(tf.pow(2., tf.cast(n, tf.float32))), tf.int32)
//  [python]
def _get_delta_filter(filt_shape, filter_support_start, filt_dtype):
  """Returns a delta filter that acts like an identity function."""
  filt_pre_shape = filt_shape[:-1]
  filter_len = filt_shape[-1]
  one_hot_loc = tf.maximum(-filter_support_start, 0)
  one_hot_loc_tensor = one_hot_loc * tf.ones(filt_pre_shape, dtype=tf.int32)
  return tf.to_float(tf.one_hot(one_hot_loc_tensor, filter_len))
//  [python]
def _get_delta_filter(filt_shape, filter_support_start, filt_dtype):
  """Returns a delta filter that acts like an identity function."""
  filt_pre_shape = filt_shape[:-1]
  filter_len = filt_shape[-1]
  one_hot_loc = tf.maximum(-filter_support_start, 0)
  one_hot_loc_tensor = one_hot_loc * tf.ones(filt_pre_shape, dtype=tf.int32)
  return tf.cast(tf.one_hot(one_hot_loc_tensor, filter_len), filt_dtype)
//  [python]
def nap_project_scalar_filter(filt: tf.Tensor,
                              near_all_pass_alpha: float = 0.7) -> tf.Tensor:
  """Limit magnitude of the complex scalar between nap_alpha and 1/nap_alpha."""
  filt_mag = tf.math.abs(filt)
  filt_phase = tf.math.angle(filt)
  filt_mag = tf.clip_by_value(filt_mag, near_all_pass_alpha,
                              1.0 / near_all_pass_alpha)
  return t.cast(filt_mag, tf.complex64) * tf.math.exp(
      1j * t.cast(filt_phase, tf.complex64))
//  [python]
def nap_project_scalar_filter(filt: tf.Tensor,
                              near_all_pass_alpha: float = 0.7) -> tf.Tensor:
  """Limit magnitude of the complex scalar between nap_alpha and 1/nap_alpha."""
  filt_mag = tf.math.abs(filt)
  filt_phase = tf.math.angle(filt)
  filt_mag = tf.clip_by_value(filt_mag, near_all_pass_alpha,
                              1.0 / near_all_pass_alpha)
  return tf.cast(filt_mag, tf.complex64) * tf.math.exp(
      1j * tf.cast(filt_phase, tf.complex64))
//  [python]
def tf_bernoulli(shape, probability):
        import tensorflow as tf

        prob_matrix = tf.fill(shape, probability)
        return tf.to_float(prob_matrix - tf.random.uniform(shape, 0, 1) >= 0)
//  [python]
def tf_bernoulli(shape, probability):
        import tensorflow as tf

        prob_matrix = tf.fill(shape, probability)
        return tf.cast(prob_matrix - tf.random.uniform(shape, 0, 1) >= 0, tf.bool)
//  [python]
def convolve(a: tf.Tensor, b: tf.Tensor) -> tf.Tensor:
  """Performs (batched) 1D convolution between two tensors.

  c[..., n] = sum_k a[..., k] b[..., n-k]

  Args:
    a: First tensor (..., seqlen_a).
    b: Second tensor (..., seqlen_b).
  Returns:
    Batched convolution result of shape (..., seqlen_a + seqlen_b - 1).
  """
  if a.dtype.is_floating and b.dtype.is_floating:
    out_dtype = a.dtype
  elif a.dtype.is_complex:
    out_dtype = a.dtype
  elif b.dtype.is_complex:
    out_dtype = b.dtype
  ab_len = _smart_dim(a, -1) + _smart_dim(b, -1) - 1
  fft_len = _enclosing_power_of_2(ab_len)
  afft = tf.signal.fft(tf.complex64(_pad_to_length(a, fft_len)))
  bfft = tf.signal.fft(tf.complex64(_pad_to_length(b, fft_len)))
  ab = tf.signal.ifft(afft * bfft)
  return tf.cast(ab[..., :ab_len], out_dtype)
//  [python]
def convolve(a: tf.Tensor, b: tf.Tensor) -> tf.Tensor:
  """Performs (batched) 1D convolution between two tensors.

  c[..., n] = sum_k a[..., k] b[..., n-k]

  Args:
    a: First tensor (..., seqlen_a).
    b: Second tensor (..., seqlen_b).
  Returns:
    Batched convolution result of shape (..., seqlen_a + seqlen_b - 1).
  """
  if a.dtype.is_floating and b.dtype.is_floating:
    out_dtype = a.dtype
  elif a.dtype.is_complex:
    out_dtype = a.dtype
  elif b.dtype.is_complex:
    out_dtype = b.dtype
  ab_len = _smart_dim(a, -1) + _smart_dim(b, -1) - 1
  fft_len = _enclosing_power_of_2(ab_len)
  afft = tf.signal.fft(tf.cast(_pad_to_length(a, fft_len), tf.complex64))
  bfft = tf.signal.fft(tf.cast(_pad_to_length(b, fft_len), tf.complex64))
  ab = tf.signal.ifft(afft * bfft)
  return tf.cast(ab[..., :ab_len], out_dtype)