jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def dense_resample(im, flow_im, output_valid_mask, name='dense_resample'):
  """ Resample reward at particular locations.
  Args:
    im:      ...xHxWxC matrix to sample from.
    flow_im: ...xHxWx2 matrix, samples the image using absolute offsets as given
             by the flow_im.
  """
  with tf.name_scope(name):
    valid_mask = None
    
    x, y = tf.unstack(flow_im, axis=-1)
    x = tf.reshape(x, [-1])
    y = tf.reshape(y, [-1])

    # constants
    shape = tf.unstack(tf.shape(im))
    channels = shape[-1]
    width = shape[-2]
    height = shape[-3]
    num_batch = tf.reduce_prod(tf.stack(shape[:-3]))
    zero = tf.constant(0, dtype=tf.int32)

    # Round up and down.
    x0 = tf.floor(x); x1 = x0 + 1;
    y0 = tf.floor(y); y1 = y0 + 1;
    
    if output_valid_mask:
      valid_mask = tf.logical_and(
          tf.logical_and(tf.less_equal(x, tf.cast(width, tf.float32)-1.), tf.greater_equal(x, 0.)),
          tf.logical_and(tf.less_equal(y, tf.cast(height, tf.float32)-1.), tf.greater_equal(y, 0.)))
      valid_mask = tf.reshape(valid_mask, shape=shape[:-1] + [1])
  
    x0 = tf.clip_by_value(x0, zero, width-1)
    x1 = tf.clip_by_value(x1, zero, width-1)
    y0 = tf.clip_by_value(y0, zero, height-1)
    y1 = tf.clip_by_value(y1, zero, height-1)

    dim2 = width; dim1 = width * height;

    # Create base index
    base = tf.reshape(tf.range(num_batch) * dim1, shape=[-1,1])
    base = tf.reshape(tf.tile(base, [1, height*width]), shape=[-1])

    base_y0 = base + y0 * dim2
    base_y1 = base + y1 * dim2
    idx_a = base_y0 + x0
    idx_b = base_y1 + x0
    idx_c = base_y0 + x1
    idx_d = base_y1 + x1

    # use indices to lookup pixels in the flat image and restore channels dim
    sh = tf.stack([tf.constant(-1,dtype=tf.int32), channels])
    im_flat = tf.reshape(im, sh)
    pixel_a = tf.gather(im_flat, idx_a)
    pixel_b = tf.gather(im_flat, idx_b)
    pixel_c = tf.gather(im_flat, idx_c)
    pixel_d = tf.gather(im_flat, idx_d)

    # and finally calculate interpolated values
    x1_f = tf.to_float(x1)
    y1_f = tf.to_float(y1)

    wa = tf.expand_dims(((x1_f - x) * (y1_f - y)), 1)
    wb = tf.expand_dims((x1_f - x) * (1.0 - (y1_f - y)), 1)
    wc = tf.expand_dims(((1.0 - (x1_f - x)) * (y1_f - y)), 1)
    wd = tf.expand_dims(((1.0 - (x1_f - x)) * (1.0 - (y1_f - y))), 1)

    output = tf.add_n([wa * pixel_a, wb * pixel_b, wc * pixel_c, wd * pixel_d])
    output = tf.reshape(output, shape=tf.shape(im))
    return output, valid_mask
//  [python]
def dense_resample(im, flow_im, output_valid_mask, name='dense_resample'):
  """ Resample reward at particular locations.
  Args:
    im:      ...xHxWxC matrix to sample from.
    flow_im: ...xHxWx2 matrix, samples the image using absolute offsets as given
             by the flow_im.
  """
  with tf.name_scope(name):
    valid_mask = None
    
    x, y = tf.unstack(flow_im, axis=-1)
    x = tf.cast(tf.reshape(x, [-1]), tf.float32)
    y = tf.cast(tf.reshape(y, [-1]), tf.float32)

    # constants
    shape = tf.unstack(tf.shape(im))
    channels = shape[-1]
    width = shape[-2]
    height = shape[-3]
    num_batch = tf.cast(tf.reduce_prod(tf.stack(shape[:-3])), 'int32')
    zero = tf.constant(0, dtype=tf.int32)

    # Round up and down.
    x0 = tf.cast(tf.floor(x), 'int32'); x1 = x0 + 1;
    y0 = tf.cast(tf.floor(y), 'int32'); y1 = y0 + 1;
    
    if output_valid_mask:
      valid_mask = tf.logical_and(
          tf.logical_and(tf.less_equal(x, tf.cast(width, tf.float32)-1.), tf.greater_equal(x, 0.)),
          tf.logical_and(tf.less_equal(y, tf.cast(height, tf.float32)-1.), tf.greater_equal(y, 0.)))
      valid_mask = tf.reshape(valid_mask, shape=shape[:-1] + [1])
  
    x0 = tf.clip_by_value(x0, zero, width-1)
    x1 = tf.clip_by_value(x1, zero, width-1)
    y0 = tf.clip_by_value(y0, zero, height-1)
    y1 = tf.clip_by_value(y1, zero, height-1)

    dim2 = width; dim1 = width * height;

    # Create base index
    base = tf.reshape(tf.range(num_batch) * dim1, shape=[-1,1])
    base = tf.reshape(tf.tile(base, [1, height*width]), shape=[-1])

    base_y0 = base + y0 * dim2
    base_y1 = base + y1 * dim2
    idx_a = base_y0 + x0
    idx_b = base_y1 + x0
    idx_c = base_y0 + x1
    idx_d = base_y1 + x1

    # use indices to lookup pixels in the flat image and restore channels dim
    sh = tf.stack([tf.constant(-1,dtype=tf.int32), channels])
    im_flat = tf.cast(tf.reshape(im, sh), dtype=tf.float32)
    pixel_a = tf.gather(im_flat, idx_a)
    pixel_b = tf.gather(im_flat, idx_b)
    pixel_c = tf.gather(im_flat, idx_c)
    pixel_d = tf.gather(im_flat, idx_d)

    # and finally calculate interpolated values
    x1_f = tf.to_float(x1)
    y1_f = tf.to_float(y1)

    wa = tf.expand_dims(((x1_f - x) * (y1_f - y)), 1)
    wb = tf.expand_dims((x1_f - x) * (1.0 - (y1_f - y)), 1)
    wc = tf.expand_dims(((1.0 - (x1_f - x)) * (y1_f - y)), 1)
    wd = tf.expand_dims(((1.0 - (x1_f - x)) * (1.0 - (y1_f - y))), 1)

    output = tf.add_n([wa * pixel_a, wb * pixel_b, wc * pixel_c, wd * pixel_d])
    output = tf.reshape(output, shape=tf.shape(im))
    return output, valid_mask
//  [python]
def lr_mask(nd, ns, dtype, *, batch=None):
    """ Left right causal mask, 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.
    :return: shape [nd, ns] or [batch, nd, ns]
    """
    i = tf.range(nd)[:, None]
    j = tf.range(ns)
    m = i >= j - ns + nd
    m = t.cast(m, dtype)
    if batch is not None:
        m = repeat(input=tf.expand_dims(m, axis=0), num=batch)
    return m
//  [python]
def lr_mask(nd, ns, dtype, *, batch=None):
    """ Left right causal mask, 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.
    :return: shape [nd, ns] or [batch, nd, ns]
    """
    i = tf.range(nd)[:, None]
    j = tf.range(ns)
    m = i >= j - ns + nd
    m = tf.cast(m, dtype)
    if batch is not None:
        m = repeat(input=tf.expand_dims(m, axis=0), num=batch)
    return m
//  [python]
def __init__(self,
                 output,
                 num_bits_per_symbol,
                 hard_out=False,
                 l=10,
                 beta=0.9,
                 dtype=tf.complex64,
                 **kwargs):
        super().__init__(dtype=dtype, **kwargs)
        assert dtype in [tf.complex64, tf.complex128], \
            "Invalid dtype"
        self._cdtype = tf.dtypes.as_dtype(dtype)
        self._rdtype = self._cdtype.real_dtype

        # Variable used to avoid numerical instabilities
        # See paragraph after Eq. (38)
        if self.dtype=="complex64":
            self._prec = 1e-6
        else:
            self._prec = 1e-12

        assert output in ("bit", "symbol"), "Unknown output"
        self._output = output

        self._hard_out = hard_out

        if self._output=="symbol":
            self._pam2qam = PAM2QAM(num_bits_per_symbol, hard_out)
        else:
            self._symbollogits2llrs = SymbolLogits2LLRs("maxlog",
                                                        num_bits_per_symbol//2,
                                                        hard_out=hard_out)
            self._demapper = Demapper("maxlog", "pam", num_bits_per_symbol//2)

        assert l>=1, "l must be a positive integer"
        self._l = l

        assert 0.0<= beta <=1.0, "beta must be in [0,1]"
        self._beta = beta

        # Create PAM constellations for real-valued detection
        self._num_bits_per_symbol = num_bits_per_symbol//2
        points = Constellation("pam", int(self._num_bits_per_symbol)).points

        # Scale constellation points to half the energy because QAM is assumed
        self._points = tf.cast(points/np.sqrt(2.0), self._rdtype)

        # Average symbol energy
        self._es = tf.constant(np.var(self._points), self._rdtype)
//  [python]
def rnnt_loss(logits, labels, label_length, logit_length, blank=0):
    if not tf.config.list_physical_devices('GPU'):
        logits = tf.nn.log_softmax(logits)
    loss = warp_rnnt_loss(
        acts=t.cast(logits, tf.float32),
        label_lengths=t.cast(label_length, tf.int32),
        labels=t.cast(labels, tf.int32),
        input_lengths=t.cast(logit_length, tf.int32),
        blank_label=blank
    )
    return loss
//  [python]
def rnnt_loss(logits, labels, label_length, logit_length, blank=0):
    if not tf.config.list_physical_devices('GPU'):
        logits = tf.nn.log_softmax(logits)
    loss = warp_rnnt_loss(
        acts=tf.cast(logits, tf.float32),
        label_lengths=tf.cast(label_length, tf.int32),
        labels=tf.cast(labels, tf.int32),
        input_lengths=tf.cast(logit_length, tf.int32),
        blank_label=blank
    )
    return loss
//  [python]
def rnnt_ctc_loss(logits, labels, label_length, logit_length, blank=0):
    logits = tf.reduce_sum(logits, 2)
    _,_,c=shape_list(logits)
    if (c-1)==blank:
        logits=tf.nn.softmax(logits,-1)
        return tf.keras.backend.ctc_batch_cost(labels,logits,tf.expand_dims(logit_length,-1),tf.expand_dims(label_length,-1))
    else:

        return tf.nn.ctc_loss(
        labels=t.cast(labels, tf.int32),
        logit_length=t.cast(logit_length, tf.int32),
        logits=t.cast(logits, tf.float32),
        label_length=t.cast(label_length, tf.int32),
        logits_time_major=False,
        blank_index=blank
    )
//  [python]
def rnnt_ctc_loss(logits, labels, label_length, logit_length, blank=0):
    logits = tf.reduce_sum(logits, 2)
    _,_,c=shape_list(logits)
    if (c-1)==blank:
        logits=tf.nn.softmax(logits,-1)
        return tf.keras.backend.ctc_batch_cost(labels,logits,tf.expand_dims(logit_length,-1),tf.expand_dims(label_length,-1))
    else:

        return tf.nn.ctc_loss(
        labels=tf.cast(labels, tf.int32),
        logit_length=tf.cast(logit_length, tf.int32),
        logits=tf.cast(logits, tf.float32),
        label_length=tf.cast(label_length, tf.int32),
        logits_time_major=False,
        blank_index=blank
    )
//  [python]
def positions_for(tokens, *, offset=None, pad_mask=None):
    """
    :param tokens: must be a 2d tensor, typically shape [batch, window]
    :param offset: scalar or rank 2 tensor ([batch, 1]) to offset the returned indices by this amount
    :param pad_mask: used for batched inputs, 0 if token is masked, 1 else. Shape [batch, window]
    :return: positions has shape tokens.shape, if pad_mask is None then
             positions[i,j] = offset + j, or
             positions[i,j] = offset[i, 0] + j
    """
    batch_size = tf.shape(tokens)[0]
    nsteps = tf.shape(tokens)[1]
    if offset is None:
        offset = 0
    positions = offset + expand_tile(tf.range(nsteps), batch_size)

    if pad_mask is None:
        return positions

    # if we have a pad_mask (i.e batched inputs), we need to shift positions
    # by the number of masked tokens. so, we count the number of masked tokens,
    # shift the positions, and then mask off the positions that aren't being used (otherwise
    # they would be negative, which doesn't work)
    n_ctx = shape_list(pad_mask)[1]
    pad_mask = tf.to_int32(pad_mask)
    # NOTE: use n_ctx - reduce_sum(mask) instead of reduce_sum(1-mask) as the sub op
    # is noticeable in profiles, esp when pad_mask is large & batched

    n_masked = n_ctx - tf.reduce_sum(pad_mask, axis=1, keepdims=True)
    return (positions - n_masked) * pad_mask
//  [python]
def positions_for(tokens, *, offset=None, pad_mask=None):
    """
    :param tokens: must be a 2d tensor, typically shape [batch, window]
    :param offset: scalar or rank 2 tensor ([batch, 1]) to offset the returned indices by this amount
    :param pad_mask: used for batched inputs, 0 if token is masked, 1 else. Shape [batch, window]
    :return: positions has shape tokens.shape, if pad_mask is None then
             positions[i,j] = offset + j, or
             positions[i,j] = offset[i, 0] + j
    """
    batch_size = tf.shape(tokens)[0]
    nsteps = tf.shape(tokens)[1]
    if offset is None:
        offset = 0
    positions = offset + expand_tile(tf.range(nsteps), batch_size)

    if pad_mask is None:
        return positions

    # if we have a pad_mask (i.e batched inputs), we need to shift positions
    # by the number of masked tokens. so, we count the number of masked tokens,
    # shift the positions, and then mask off the positions that aren't being used (otherwise
    # they would be negative, which doesn't work)
    n_ctx = shape_list(pad_mask)[1]
    pad_mask = tf.cast(pad_mask, tf.int32)
    # NOTE: use n_ctx - reduce_sum(mask) instead of reduce_sum(1-mask) as the sub op
    # is noticeable in profiles, esp when pad_mask is large & batched

    n_masked = n_ctx - tf.reduce_sum(pad_mask, axis=1, keepdims=True)
    return (positions - n_masked) * pad_mask
//  [python]
def _create_position_ids_from_input_ids(
    input_ids: tf.Tensor, past_key_values_length: int, padding_idx: Optional[int]
) -> tf.Tensor:
    """
    Replace non-padding symbols with their position numbers. Position numbers begin at padding_idx+1. Padding symbols
    are ignored. This is modified from fairseq's `utils.make_positions`.
    """
    # The series of casts and type-conversions here are carefully balanced to both work with ONNX export and XLA.
    mask = tf.where(input_ids != padding_idx, 1, 0)
    incremental_indices = (t.cast(tf.cumsum(mask, axis=1), dtype=mask.dtype) + past_key_values_length) * mask
    return t.cast(incremental_indices, dtype=tf.int64) + padding_idx
//  [python]
def _create_position_ids_from_input_ids(
    input_ids: tf.Tensor, past_key_values_length: int, padding_idx: Optional[int]
) -> tf.Tensor:
    """
    Replace non-padding symbols with their position numbers. Position numbers begin at padding_idx+1. Padding symbols
    are ignored. This is modified from fairseq's `utils.make_positions`.
    """
    # The series of casts and type-conversions here are carefully balanced to both work with ONNX export and XLA.
    mask = tf.where(input_ids != padding_idx, 1, 0)
    incremental_indices = (tf.cast(tf.cumsum(mask, axis=1), dtype=mask.dtype) + past_key_values_length) * mask
    return tf.cast(incremental_indices, dtype=tf.int64) + padding_idx
//  [python]
def single_call(self, arg):
    return t.cast(arg, self._dtype)
//  [python]
def single_call(self, arg):
    return tf.cast(arg, self._dtype)
//  [python]
def __call__(self, step):
    with tf.name_scope(self.name or 'WarmUp') as name:
      # Implements polynomial warmup. i.e., if global_step < warmup_steps, the
      # learning rate will be `global_step/num_warmup_steps * init_lr`.
      global_step_float = t.cast(step, tf.float32)
      warmup_steps_float = t.cast(self.warmup_steps, tf.float32)
      warmup_percent_done = global_step_float / warmup_steps_float
      warmup_learning_rate = (
          self.initial_learning_rate *
          tf.math.pow(warmup_percent_done, self.power))
      return tf.cond(
          global_step_float < warmup_steps_float,
          lambda: warmup_learning_rate,
          lambda: self.decay_schedule_fn(step),
          name=name)
//  [python]
def __call__(self, step):
    with tf.name_scope(self.name or 'WarmUp') as name:
      # Implements polynomial warmup. i.e., if global_step < warmup_steps, the
      # learning rate will be `global_step/num_warmup_steps * init_lr`.
      global_step_float = tf.cast(step, tf.float32)
      warmup_steps_float = tf.cast(self.warmup_steps, tf.float32)
      warmup_percent_done = global_step_float / warmup_steps_float
      warmup_learning_rate = (
          self.initial_learning_rate *
          tf.math.pow(warmup_percent_done, self.power))
      return tf.cond(
          global_step_float < warmup_steps_float,
          lambda: warmup_learning_rate,
          lambda: self.decay_schedule_fn(step),
          name=name)
//  [python]
def single_call(self, arg):
    if arg.dtype == tf.string:
      arg = tf.reshape(arg, ())
      return tf.strings.length(arg)
    else:
      mask = tf.logical_and(self._vocab.padding_mask(arg),
                            self._vocab.special_token_mask(arg))
      return tf.reduce_sum(t.cast(mask, tf.int32))
//  [python]
def single_call(self, arg):
    if arg.dtype == tf.string:
      arg = tf.reshape(arg, ())
      return tf.strings.length(arg)
    else:
      mask = tf.logical_and(self._vocab.padding_mask(arg),
                            self._vocab.special_token_mask(arg))
      return tf.reduce_sum(tf.cast(mask, tf.int32))
//  [python]
def single_call(self, arg):
    """Expects positions to be a tf.Tensor<float>[n, 3]."""
    # Makes a batch of size 1 to be compatible with pairwise_square_dist.
    pos = tf.expand_dims(arg, 0)
    sq_dist = pairs.square_distances(pos, pos)[0]
    return t.cast(sq_dist < self._threshold ** 2, dtype=tf.float32)
//  [python]
def single_call(self, arg):
    """Expects positions to be a tf.Tensor<float>[n, 3]."""
    # Makes a batch of size 1 to be compatible with pairwise_square_dist.
    pos = tf.expand_dims(arg, 0)
    sq_dist = pairs.square_distances(pos, pos)[0]
    return tf.cast(sq_dist < self._threshold ** 2, dtype=tf.float32)
//  [python]
def single_call(self, arg):
    """For later."""
    # Makes a batch of size 1 to be compatible with pairwise_square_dist.
    arg = t.cast(arg[None, :] * arg[:, None], dtype=tf.float32)
    arg.set_shape(self._shape)
    return arg
//  [python]
def single_call(self, arg):
    """For later."""
    # Makes a batch of size 1 to be compatible with pairwise_square_dist.
    arg = tf.cast(arg[None, :] * arg[:, None], dtype=tf.float32)
    arg.set_shape(self._shape)
    return arg
//  [python]
def _get_len(self, t):
    if not self._precomputed:
      if t.dtype == tf.string:
        t = tf.reshape(t, ())
        return tf.strings.length(t)
      else:
        mask = tf.logical_and(
            self._vocab.padding_mask(t), self._vocab.special_token_mask(t))
        return tf.reduce_sum(t.astype(tf.int32))
    return t
//  [python]
def _get_len(self, t):
    if not self._precomputed:
      if t.dtype == tf.string:
        t = tf.reshape(t, ())
        return tf.strings.length(t)
      else:
        mask = tf.logical_and(
            self._vocab.padding_mask(t), self._vocab.special_token_mask(t))
        return tf.reduce_sum(tf.cast(mask, tf.int32))
    return t
//  [python]
def call(self, inputs):

        # Flatten the batch dimensions
        y, h, s = inputs
        batch_shape = tf.shape(y)[:-1]
        num_batch_dims = len(batch_shape)
        if num_batch_dims > 1:
            y = flatten_dims(y, num_batch_dims, 0)
            h = flatten_dims(h, num_batch_dims, 0)
            s = flatten_dims(s, num_batch_dims, 0)
            inputs = (y,h,s)

        # Number of transmit streams
        n_t = tf.shape(h)[-1]

        # Whiten channel
        y, h, s = whiten_channel(y, h, s)

        # Convert channel to real-valued representation
        y, h, s = complex2real_channel(y,h,s)

        # Convert all inputs to desired dtypes
        y = t.cast(y, self._rdtype)
        h = t.cast(h, self._rdtype)
        no = t.cast(0.5, self._rdtype)

        # Gather relevant parameters
        batch_dims = tf.shape(y)[:-1]
        n_t_r = tf.shape(h)[-1]

        # Initialize gamma and lambda (Paragraph after Eq. (29))
        gam = tf.zeros(tf.concat([batch_dims, [n_t_r]], axis=0), y.dtype)
        lam = tf.ones(tf.concat([batch_dims, [n_t_r]], axis=0), y.dtype)
        lam /= t.cast(self._es, y.dtype)

        # Precompute values that are repeatedly needed
        h_t_h = tf.matmul(h, h, transpose_a=True)
        y = tf.expand_dims(y, axis=-1)
        h_t_y = tf.matmul(h, y, transpose_a=True)
        no = expand_to_rank(no, tf.rank(h), axis=-1)

        for _ in range(self._l):
            sigma, mu = self.compute_sigma_mu(h_t_h, h_t_y, no, lam, gam)
            v_obs, x_obs = self.compute_v_x_obs(sigma, mu, lam, gam)
            v, x, logits = self.compute_v_x(v_obs, x_obs)
            lam, gam = self.update_lam_gam(v, v_obs, x, x_obs, lam, gam)

        # Extract the logits for the 2 PAM constellations for each streams
        pam1_logits = logits[...,:n_t,:]
        pam2_logits = logits[...,n_t:,:]

        if self._output=="symbol" and self._hard_out:
            # Take hard decisions on PAM symbol;s
            pam1_ind = tf.argmax(pam1_logits, axis=-1, output_type=tf.int32)
            pam2_ind = tf.argmax(pam2_logits, axis=-1, output_type=tf.int32)

            # Transform to QAM indices
            qam_ind = self._pam2qam(pam1_ind, pam2_ind)

            # Reshape batch dimensions
            if num_batch_dims > 1:
                qam_ind = split_dim(qam_ind, batch_shape, 0)

            return qam_ind

        elif self._output=="symbol" and not self._hard_out:
            qam_logits = self._pam2qam(pam1_logits, pam2_logits)

            # Reshape batch dimensions
            if num_batch_dims > 1:
                qam_logits = split_dim(qam_logits, batch_shape, 0)

            return qam_logits

        elif self._output=="bit":
            # Compute LLRs for both PAM constellations
            llr1 = self._symbollogits2llrs(pam1_logits)
            llr2 = self._symbollogits2llrs(pam2_logits)

            # Put LLRs in the correct order and shape
            llr = tf.stack([llr1, llr2], -1)
            llr = flatten_last_dims(llr)

            # Reshape batch dimensions
            if num_batch_dims > 1:
                llr = split_dim(llr, batch_shape, 0)

            return llr
//  [python]
def call(self, inputs):

        # Flatten the batch dimensions
        y, h, s = inputs
        batch_shape = tf.shape(y)[:-1]
        num_batch_dims = len(batch_shape)
        if num_batch_dims > 1:
            y = flatten_dims(y, num_batch_dims, 0)
            h = flatten_dims(h, num_batch_dims, 0)
            s = flatten_dims(s, num_batch_dims, 0)
            inputs = (y,h,s)

        # Number of transmit streams
        n_t = tf.shape(h)[-1]

        # Whiten channel
        y, h, s = whiten_channel(y, h, s)

        # Convert channel to real-valued representation
        y, h, s = complex2real_channel(y,h,s)

        # Convert all inputs to desired dtypes
        y = tf.cast(y, self._rdtype)
        h = tf.cast(h, self._rdtype)
        no = tf.cast(0.5, self._rdtype)

        # Gather relevant parameters
        batch_dims = tf.shape(y)[:-1]
        n_t_r = tf.shape(h)[-1]

        # Initialize gamma and lambda (Paragraph after Eq. (29))
        gam = tf.zeros(tf.concat([batch_dims, [n_t_r]], axis=0), y.dtype)
        lam = tf.ones(tf.concat([batch_dims, [n_t_r]], axis=0), y.dtype)
        lam /= tf.cast(self._es, y.dtype)

        # Precompute values that are repeatedly needed
        h_t_h = tf.matmul(h, h, transpose_a=True)
        y = tf.expand_dims(y, axis=-1)
        h_t_y = tf.matmul(h, y, transpose_a=True)
        no = expand_to_rank(no, tf.rank(h), axis=-1)

        for _ in range(self._l):
            sigma, mu = self.compute_sigma_mu(h_t_h, h_t_y, no, lam, gam)
            v_obs, x_obs = self.compute_v_x_obs(sigma, mu, lam, gam)
            v, x, logits = self.compute_v_x(v_obs, x_obs)
            lam, gam = self.update_lam_gam(v, v_obs, x, x_obs, lam, gam)

        # Extract the logits for the 2 PAM constellations for each streams
        pam1_logits = logits[...,:n_t,:]
        pam2_logits = logits[...,n_t:,:]

        if self._output=="symbol" and self._hard_out:
            # Take hard decisions on PAM symbol;s
            pam1_ind = tf.argmax(pam1_logits, axis=-1, output_type=tf.int32)
            pam2_ind = tf.argmax(pam2_logits, axis=-1, output_type=tf.int32)

            # Transform to QAM indices
            qam_ind = self._pam2qam(pam1_ind, pam2_ind)

            # Reshape batch dimensions
            if num_batch_dims > 1:
                qam_ind = split_dim(qam_ind, batch_shape, 0)

            return qam_ind

        elif self._output=="symbol" and not self._hard_out:
            qam_logits = self._pam2qam(pam1_logits, pam2_logits)

            # Reshape batch dimensions
            if num_batch_dims > 1:
                qam_logits = split_dim(qam_logits, batch_shape, 0)

            return qam_logits

        elif self._output=="bit":
            # Compute LLRs for both PAM constellations
            llr1 = self._symbollogits2llrs(pam1_logits)
            llr2 = self._symbollogits2llrs(pam2_logits)

            # Put LLRs in the correct order and shape
            llr = tf.stack([llr1, llr2], -1)
            llr = flatten_last_dims(llr)

            # Reshape batch dimensions
            if num_batch_dims > 1:
                llr = split_dim(llr, batch_shape, 0)

            return llr
//  [python]
def K_diag(self, X: TensorType) -> tf.Tensor:
        cs(X, "[batch..., N, 1]  # The `Coregion` kernel requires a 1D input space.")

        X = t.cast(X[..., 0], tf.int32)
        B_diag = self.output_variance()
        return tf.gather(B_diag, X)
//  [python]
def K_diag(self, X: TensorType) -> tf.Tensor:
        cs(X, "[batch..., N, 1]  # The `Coregion` kernel requires a 1D input space.")

        X = tf.cast(X[..., 0], tf.int32)
        B_diag = self.output_variance()
        return tf.gather(B_diag, X)
//  [python]
def _undiscretize_struct(struct, step_size, tf_dtype_struct):
  """Unscales the discretized structure and casts back to original dtypes."""

  def undiscretize_tensor(x, original_dtype):
    unscaled_x = t.cast(x, original_dtype) * t.cast(step_size, original_dtype)
    return unscaled_x

  return tf.nest.map_structure(undiscretize_tensor, struct, tf_dtype_struct)
//  [python]
def _undiscretize_struct(struct, step_size, tf_dtype_struct):
  """Unscales the discretized structure and casts back to original dtypes."""

  def undiscretize_tensor(x, original_dtype):
    unscaled_x = tf.cast(x, original_dtype) * tf.cast(step_size, original_dtype)
    return unscaled_x

  return tf.nest.map_structure(undiscretize_tensor, struct, tf_dtype_struct)
//  [python]
def _discretize_struct(struct, step_size):
  """Scales and rounds each tensor of the structure to the integer grid."""

  def discretize_tensor(x):
    seed = t.cast(
        t.stack([t.timestamp() * 1e6, t.timestamp() * 1e6]), dtype=t.int64
    )
    scaled_x = t.divide(x, t.cast(step_size, x.dtype))
    prob_x = scaled_x - t.floor(scaled_x)
    random_x = t.random.stateless_uniform(x.shape, seed=seed, dtype=x.dtype)
    discretized_x = t.where(
        t.less_equal(random_x, prob_x),
        t.math.ceil(scaled_x),
        t.math.floor(scaled_x),
    )
    return t.cast(discretized_x, t.int32)

  return t.nest.map_structure(discretize_tensor, struct)
//  [python]
def _discretize_struct(struct, step_size):
  """Scales and rounds each tensor of the structure to the integer grid."""

  def discretize_tensor(x):
    seed = tf.cast(
        tf.stack([tf.timestamp() * 1e6, tf.timestamp() * 1e6]), dtype=tf.int64
    )
    scaled_x = tf.divide(x, tf.cast(step_size, x.dtype))
    prob_x = scaled_x - tf.floor(scaled_x)
    random_x = tf.random.stateless_uniform(x.shape, seed=seed, dtype=x.dtype)
    discretized_x = tf.where(
        tf.less_equal(random_x, prob_x),
        tf.math.ceil(scaled_x),
        tf.math.floor(scaled_x),
    )
    return tf.cast(discretized_x, tf.int32)

  return tf.nest.map_structure(discretize_tensor, struct)
//  [python]
def K(self, X: TensorType, X2: Optional[TensorType] = None) -> tf.Tensor:
        cs(X, "[batch..., N, 1]  # The `Coregion` kernel requires a 1D input space.")

        B = cs(self.output_covariance(), "[O, O]")
        X = cs(t.cast(X[..., 0], tf.int32), "[batch..., N]")
        if X2 is None:
            batch = tf.shape(X)[:-1]
            N = tf.shape(X)[-1]
            O = tf.shape(B)[-1]

            result = cs(tf.gather(B, X), "[batch..., N, O]")
            result = cs(tf.reshape(result, [-1, N, O]), "[flat_batch, N, O]")
            flat_X = cs(tf.reshape(X, [-1, N]), "[flat_batch, N]")
            result = cs(tf.gather(result, flat_X, axis=2, batch_dims=1), "[flat_batch, N, N]")
            result = cs(tf.reshape(result, tf.concat([batch, [N, N]], 0)), "[batch..., N, N]")
        else:
            X2 = cs(t.cast(X2[..., 0], tf.int32), "[batch2..., N2]")

            rank2 = tf.rank(X2)

            result = cs(tf.gather(B, X2), "[batch2..., N2, O]")
            result = cs(
                tf.transpose(result, tf.concat([[rank2], tf.range(rank2)], 0)), "[O, batch2..., N2]"
            )
            result = cs(tf.gather(result, X), "[batch..., N, batch2..., N2]")

        return result
//  [python]
def K(self, X: TensorType, X2: Optional[TensorType] = None) -> tf.Tensor:
        cs(X, "[batch..., N, 1]  # The `Coregion` kernel requires a 1D input space.")

        B = cs(self.output_covariance(), "[O, O]")
        X = cs(tf.cast(X[..., 0], tf.int32), "[batch..., N]")
        if X2 is None:
            batch = tf.shape(X)[:-1]
            N = tf.shape(X)[-1]
            O = tf.shape(B)[-1]

            result = cs(tf.gather(B, X), "[batch..., N, O]")
            result = cs(tf.reshape(result, [-1, N, O]), "[flat_batch, N, O]")
            flat_X = cs(tf.reshape(X, [-1, N]), "[flat_batch, N]")
            result = cs(tf.gather(result, flat_X, axis=2, batch_dims=1), "[flat_batch, N, N]")
            result = cs(tf.reshape(result, tf.concat([batch, [N, N]], 0)), "[batch..., N, N]")
        else:
            X2 = cs(tf.cast(X2[..., 0], tf.int32), "[batch2..., N2]")

            rank2 = tf.rank(X2)

            result = cs(tf.gather(B, X2), "[batch2..., N2, O]")
            result = cs(
                tf.transpose(result, tf.concat([[rank2], tf.range(rank2)], 0)), "[O, batch2..., N2]"
            )
            result = cs(tf.gather(result, X), "[batch..., N, batch2..., N2]")

        return result
//  [python]
def call(self, inputs, length=None):
    """Implements call() for the layer.

    Args:
      inputs: An tensor whose second dimension will be used as `length`. If
        `None`, the other `length` argument must be specified.
      length: An optional integer specifying the number of positions. If both
        `inputs` and `length` are spcified, `length` must be equal to the second
        dimension of `inputs`.

    Returns:
      A tensor in shape of `(length, hidden_size)`.
    """
    if inputs is None and length is None:
      raise ValueError("If inputs is None, `length` must be set in "
                       "RelativePositionEmbedding().")
    if inputs is not None:
      input_shape = tf_utils.get_shape_list(inputs)
      if length is not None and length != input_shape[1]:
        raise ValueError(
            "If inputs is not None, `length` must equal to input_shape[1].")
      length = input_shape[1]
    position = t.cast(t.range(length), t.float32)
    num_timescales = self._hidden_size // 2
    min_timescale, max_timescale = self._min_timescale, self._max_timescale
    log_timescale_increment = (
        math.log(float(max_timescale) / float(min_timescale)) /
        (t.cast(num_timescales, t.float32) - 1))
    inv_timescales = min_timescale * t.exp(
        t.cast(t.range(num_timescales), t.float32) *
        -log_timescale_increment)
    scaled_time = t.expand_dims(position, 1) * t.expand_dims(
        inv_timescales, 0)
    position_embeddings = t.concat(
        [t.sin(scaled_time), t.cos(scaled_time)], axis=1)
    return position_embeddings
//  [python]
def call(self, inputs, length=None):
    """Implements call() for the layer.

    Args:
      inputs: An tensor whose second dimension will be used as `length`. If
        `None`, the other `length` argument must be specified.
      length: An optional integer specifying the number of positions. If both
        `inputs` and `length` are spcified, `length` must be equal to the second
        dimension of `inputs`.

    Returns:
      A tensor in shape of `(length, hidden_size)`.
    """
    if inputs is None and length is None:
      raise ValueError("If inputs is None, `length` must be set in "
                       "RelativePositionEmbedding().")
    if inputs is not None:
      input_shape = tf_utils.get_shape_list(inputs)
      if length is not None and length != input_shape[1]:
        raise ValueError(
            "If inputs is not None, `length` must equal to input_shape[1].")
      length = input_shape[1]
    position = tf.cast(tf.range(length), tf.float32)
    num_timescales = self._hidden_size // 2
    min_timescale, max_timescale = self._min_timescale, self._max_timescale
    log_timescale_increment = (
        math.log(float(max_timescale) / float(min_timescale)) /
        (tf.cast(num_timescales, tf.float32) - 1))
    inv_timescales = min_timescale * tf.exp(
        tf.cast(tf.range(num_timescales), tf.float32) *
        -log_timescale_increment)
    scaled_time = tf.expand_dims(position, 1) * tf.expand_dims(
        inv_timescales, 0)
    position_embeddings = tf.concat(
        [tf.sin(scaled_time), tf.cos(scaled_time)], axis=1)
    return position_embeddings
//  [python]
def _relative_position_bucket(relative_position,
                              bidirectional=True,
                              num_buckets=32,
                              max_distance=128):
  """Translate relative position to a bucket number for relative attention.

  The relative position is defined as memory_position - query_position, i.e.
  the distance in tokens from the attending position to the attended-to
  position.

  If `bidirectional=False`, then positive relative positions are invalid.

  We use smaller buckets for small absolute relative_position and larger
  buckets for larger absolute relative_positions.

  All relative positions >=max_distance map to the same bucket.

  All relative positions <=-max_distance map to the same bucket.

  This should allow for more graceful generalization to longer sequences
  than the model has been trained on.

  Args:
    relative_position: An int32 Tensor
    bidirectional: A boolean - whether the attention is bidirectional
    num_buckets: An integer
    max_distance: An integer

  Returns:
    A Tensor with the same shape as relative_position, containing int32
    values in the range [0, num_buckets)
  """
  ret = 0
  n = -relative_position
  if bidirectional:
    num_buckets //= 2
    ret += t.cast(t.less(n, 0), t.int32) * num_buckets
    n = t.abs(n)
  else:
    n = t.maximum(n, 0)
  # now n is in the range [0, inf)
  max_exact = num_buckets // 2
  is_small = t.less(n, max_exact)
  val_if_large = max_exact + t.cast(
      t.log(t.cast(n, t.float32) / max_exact) /
      math.log(max_distance / max_exact) * (num_buckets - max_exact),
      t.int32,
  )
  val_if_large = t.minimum(val_if_large, num_buckets - 1)
  ret += t.where(is_small, n, val_if_large)
  return ret
//  [python]
def _relative_position_bucket(relative_position,
                              bidirectional=True,
                              num_buckets=32,
                              max_distance=128):
  """Translate relative position to a bucket number for relative attention.

  The relative position is defined as memory_position - query_position, i.e.
  the distance in tokens from the attending position to the attended-to
  position.

  If `bidirectional=False`, then positive relative positions are invalid.

  We use smaller buckets for small absolute relative_position and larger
  buckets for larger absolute relative_positions.

  All relative positions >=max_distance map to the same bucket.

  All relative positions <=-max_distance map to the same bucket.

  This should allow for more graceful generalization to longer sequences
  than the model has been trained on.

  Args:
    relative_position: An int32 Tensor
    bidirectional: A boolean - whether the attention is bidirectional
    num_buckets: An integer
    max_distance: An integer

  Returns:
    A Tensor with the same shape as relative_position, containing int32
    values in the range [0, num_buckets)
  """
  ret = 0
  n = -relative_position
  if bidirectional:
    num_buckets //= 2
    ret += tf.cast(tf.math.less(n, 0), tf.int32) * num_buckets
    n = tf.math.abs(n)
  else:
    n = tf.math.maximum(n, 0)
  # now n is in the range [0, inf)
  max_exact = num_buckets // 2
  is_small = tf.math.less(n, max_exact)
  val_if_large = max_exact + tf.dtypes.cast(
      tf.math.log(tf.cast(n, tf.float32) / max_exact) /
      math.log(max_distance / max_exact) * (num_buckets - max_exact),
      tf.int32,
  )
  val_if_large = tf.math.minimum(val_if_large, num_buckets - 1)
  ret += tf.where(is_small, n, val_if_large)
  return ret
//  [python]
def _get_positive_similarity_statistics(
      self,
      logits: tf.Tensor,
      inst_mask: tf.Tensor) -> Mapping[str, tf.Tensor]:
    """Gets positive examples similarity statistics.

    Args:
      logits: the logits matrix.
      inst_mask: the instance validity mask.

    Returns:
      logs: a dictionary of logs.
    """
    # logits in shape [b, n]
    # inst_mask in shape [b, n, 1]
    inst_mask = tf.squeeze(inst_mask, axis=-1)
    inst_mask = tf.cast(inst_mask, dtype=logits.dtype)
    similarity = logits * inst_mask * self._temperature

    num_instances = tf.reduce_sum(inst_mask)
    similarity_mean = tf.reduce_sum(similarity) / num_instances

    similarity_max = similarity - (1.0 - inst_mask) * _LARGE_NUM
    similarity_max = tf.reduce_max(similarity_max)

    similarity_min = similarity + (1.0 - inst_mask) * _LARGE_NUM
    similarity_min = tf.reduce_min(similarity_min)

    logs = {
        'positive_similarity_mean': similarity_mean,
        'positive_similarity_min': similarity_min,
        'positive_similarity_max': similarity_max,
    }
    return logs
//  [python]
def _create_padding_mask(x):
  mask = tf.cast(tf.math.equal(x, 0), tf.float32)
  # (batch_size, 1, 1, sequence length)
  return mask[:, tf.newaxis, tf.newaxis, :]
//  [python]
def _get_negative_similarity_statistics(
      self,
      logits: tf.Tensor,
      batch_masks: tf.Tensor,
      inst_mask: tf.Tensor) -> Mapping[str, tf.Tensor]:
    """Gets negative examples similarity statistics.

    Args:
      logits: the logits matrix.
      batch_masks: the batch validity mask.
      inst_mask: the instance validity mask.

    Returns:
      logs: a dictionary of logs.
    """
    # logits = [b, n, bl, n]
    # batch_masks = [b, n, bl, n]
    # inst_mask = [b, n, 1]
    inst_mask = t.cast(inst_mask, logits.dtype)
    batch_masks = t.cast(batch_masks, logits.dtype)
    batch_masks = tf.ones_like(batch_masks) - batch_masks
    masks = batch_masks * inst_mask[..., None]
    # Recover the raw similarity and mask self-similarity, which will be
    # removed from negative samples.
    similarity = logits * masks * self._temperature
    similarity_mean = tf.reduce_sum(similarity) / tf.reduce_sum(masks)

    similarity_masks = tf.squeeze(inst_mask, axis=-1)
    similarity_max = similarity - (1.0 - masks) * _LARGE_NUM
    similarity_max = tf.reduce_max(similarity_max, axis=[-1, -2])
    similarity_max = tf.reduce_sum(
        similarity_max * similarity_masks) / tf.reduce_sum(similarity_masks)

    similarity_min = similarity + (1.0 - masks) * _LARGE_NUM
    similarity_min = tf.reduce_min(similarity_min, axis=[-1, -2])
    similarity_min = tf.reduce_sum(
        similarity_min * similarity_masks) / tf.reduce_sum(similarity_masks)
    logs = {
        'negative_similarity_mean': similarity_mean,
        'negative_similarity_min': similarity_min,
        'negative_similarity_max': similarity_max,
    }
    return logs
//  [python]
def _get_negative_similarity_statistics(
      self,
      logits: tf.Tensor,
      batch_masks: tf.Tensor,
      inst_mask: tf.Tensor) -> Mapping[str, tf.Tensor]:
    """Gets negative examples similarity statistics.

    Args:
      logits: the logits matrix.
      batch_masks: the batch validity mask.
      inst_mask: the instance validity mask.

    Returns:
      logs: a dictionary of logs.
    """
    # logits = [b, n, bl, n]
    # batch_masks = [b, n, bl, n]
    # inst_mask = [b, n, 1]
    inst_mask = tf.cast(inst_mask, logits.dtype)
    batch_masks = tf.cast(batch_masks, logits.dtype)
    batch_masks = tf.ones_like(batch_masks) - batch_masks
    masks = batch_masks * inst_mask[..., None]
    # Recover the raw similarity and mask self-similarity, which will be
    # removed from negative samples.
    similarity = logits * masks * self._temperature
    similarity_mean = tf.reduce_sum(similarity) / tf.reduce_sum(masks)

    similarity_masks = tf.squeeze(inst_mask, axis=-1)
    similarity_max = similarity - (1.0 - masks) * _LARGE_NUM
    similarity_max = tf.reduce_max(similarity_max, axis=[-1, -2])
    similarity_max = tf.reduce_sum(
        similarity_max * similarity_masks) / tf.reduce_sum(similarity_masks)

    similarity_min = similarity + (1.0 - masks) * _LARGE_NUM
    similarity_min = tf.reduce_min(similarity_min, axis=[-1, -2])
    similarity_min = tf.reduce_sum(
        similarity_min * similarity_masks) / tf.reduce_sum(similarity_masks)
    logs = {
        'negative_similarity_mean': similarity_mean,
        'negative_similarity_min': similarity_min,
        'negative_similarity_max': similarity_max,
    }
    return logs
//  [python]
def get_angles(self, position, i, d_model):
    angles = 1 / tf.pow(10000, (2 * (i // 2)) / t.float32(d_model))
    return position * angles
//  [python]
def get_angles(self, position, i, d_model):
    angles = 1 / tf.pow(10000, (2 * (i // 2)) / tf.cast(d_model, tf.float32))
    return position * angles
//  [python]
def _scaled_dot_product_attention(query, key, value, mask):
  """Actual attention function using dot product."""
  matmul_qk = tf.matmul(query, key, transpose_b=True)
  depth = tf.cast(tf.shape(key)[-1], tf.float32)
  logits = matmul_qk / tf.math.sqrt(depth)

  # add the mask zero out padding tokens.
  if mask is not None:
    logits += (mask * -1e9)

  attention_weights = tf.nn.softmax(logits, axis=-1)
  return tf.matmul(attention_weights, value), attention_weights
//  [python]
def positional_encoding(self, position, d_model):
    angle_rads = self.get_angles(
        position=tf.range(position, dtype=tf.float32)[:, tf.newaxis],
        i=tf.range(d_model, dtype=tf.float32)[tf.newaxis, :],
        d_model=d_model)
    # apply sin to even index in the array
    sines = tf.math.sin(angle_rads[:, 0::2])
    # apply cos to odd index in the array
    cosines = tf.math.cos(angle_rads[:, 1::2])

    pos_encoding = tf.concat([sines, cosines], axis=-1)
    pos_encoding = pos_encoding[tf.newaxis, Ellipsis]
    return tf.to_float(pos_encoding)
//  [python]
def positional_encoding(self, position, d_model):
    angle_rads = self.get_angles(
        position=tf.range(position, dtype=tf.float32)[:, tf.newaxis],
        i=tf.range(d_model, dtype=tf.float32)[tf.newaxis, :],
        d_model=d_model)
    # apply sin to even index in the array
    sines = tf.math.sin(angle_rads[:, 0::2])
    # apply cos to odd index in the array
    cosines = tf.math.cos(angle_rads[:, 1::2])

    pos_encoding = tf.concat([sines, cosines], axis=-1)
    pos_encoding = pos_encoding[tf.newaxis, Ellipsis]
    return tf.cast(pos_encoding, tf.float32)
//  [python]
def __init__(self, d_model, warmup_steps=4000):
    super(CustomSchedule, self).__init__()
    self.d_model = d_model
    self.d_model = tf.cast(self.d_model, tf.float32)
    self.warmup_steps = warmup_steps
//  [python]
def __call__(self,
               inputs: tf.Tensor,
               num_replicas: int = 1) -> Mapping[str, tf.Tensor]:
    """Calculates the loss.

    Args:
      inputs: the embeddings (in shape [2*B, C]) from video clips after the
        projection head.
      num_replicas: the number of TPU replicas.

    Returns:
      a dictionary contains calculated loss and statistics.
    """
    inputs1, inputs2 = tf.split(inputs, num_or_size_splits=2, axis=0)
    if self._normalize_inputs:
      inputs1 = tf.math.l2_normalize(inputs1, -1)
      inputs2 = tf.math.l2_normalize(inputs2, -1)
    batch_size = tf.shape(inputs1)[0]

    if num_replicas == 1:
      # This is the local version.
      inputs1_large = inputs1
      inputs2_large = inputs2
      labels = tf.one_hot(tf.range(batch_size), batch_size * 2)
      masks = tf.one_hot(tf.range(batch_size), batch_size)
    else:
      # This is the cross-tpu version.
      inputs1_large = tpu_cross_replica_concat(inputs1, num_replicas)
      inputs2_large = tpu_cross_replica_concat(inputs2, num_replicas)
      enlarged_batch_size = tf.shape(inputs1_large)[0]
      replica_id = tf.cast(tf.cast(xla.replica_id(), tf.uint32), tf.int32)
      labels_idx = tf.range(batch_size) + replica_id * batch_size
      labels = tf.one_hot(labels_idx, enlarged_batch_size * 2)
      masks = tf.one_hot(labels_idx, enlarged_batch_size)

    logits_aa = tf.matmul(
        inputs1, inputs1_large, transpose_b=True) / self._temperature
    logits_aa = logits_aa - tf.cast(masks, logits_aa.dtype) * _LARGE_NUM
    logits_bb = tf.matmul(
        inputs2, inputs2_large, transpose_b=True) / self._temperature
    logits_bb = logits_bb - tf.cast(masks, logits_bb.dtype) * _LARGE_NUM
    logits_ab = tf.matmul(
        inputs1, inputs2_large, transpose_b=True) / self._temperature
    logits_ba = tf.matmul(
        inputs2, inputs1_large, transpose_b=True) / self._temperature

    loss_a = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(
        labels, tf.concat([logits_ab, logits_aa], 1)))
    loss_b = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(
        labels, tf.concat([logits_ba, logits_bb], 1)))
    loss = loss_a + loss_b

    contrast_prob = tf.nn.softmax(logits_ab)
    contrast_entropy = - tf.reduce_mean(
        tf.reduce_sum(contrast_prob * tf.math.log(contrast_prob + 1e-8), -1))

    contrast_acc = tf.equal(tf.argmax(labels, 1), tf.argmax(logits_ab, axis=1))
    contrast_acc = tf.reduce_mean(tf.cast(contrast_acc, tf.float32))

    return {
        'loss': loss,
        'contrastive_accuracy': contrast_acc,
        'contrastive_entropy': contrast_entropy,
    }
//  [python]
def _encoder(vocab_size,
             num_layers,
             units,
             d_model,
             num_heads,
             dropout,
             name="encoder"):
  """Encoder component."""
  inputs = tf.keras.Input(shape=(None,), name="inputs")
  padding_mask = tf.keras.Input(shape=(1, 1, None), name="padding_mask")

  embeddings = tf.keras.layers.Embedding(vocab_size, d_model)(inputs)
  embeddings *= tf.math.sqrt(tf.cast(d_model, tf.float32))
  embeddings = PositionalEncoding(_MAX_SEQUENCE_LENGTH, d_model)(embeddings)
  outputs = tf.keras.layers.Dropout(rate=dropout)(embeddings)

  for i in range(num_layers):
    outputs = _encoder_layer(
        units=units,
        d_model=d_model,
        num_heads=num_heads,
        dropout=dropout,
        name="encoder_layer_{}".format(i),
    )([outputs, padding_mask])

  return tf.keras.Model(
      inputs=[inputs, padding_mask], outputs=outputs, name=name)