jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def _mask_hidden_states(self, hidden_states: tf.Tensor, mask_time_indices: tf.Tensor | None = None):
        """
        Masks extracted features along time axis and/or along feature axis according to
        [SpecAugment](https://arxiv.org/abs/1904.08779).
        """
        batch_size, sequence_length, hidden_size = shape_list(hidden_states)

        # `config.apply_spec_augment` can set masking to False
        if not getattr(self.config, "apply_spec_augment", True):
            return hidden_states

        if mask_time_indices is not None:
            # apply SpecAugment along time axis with given mask_time_indices
            hidden_states = tf.where(
                tf.cast(mask_time_indices[:, :, tf.newaxis], tf.bool),
                self.masked_spec_embed[tf.newaxis, tf.newaxis, :],
                hidden_states,
            )

        elif self.config.mask_time_prob > 0:
            # generate indices & apply SpecAugment along time axis
            mask_time_indices = _compute_mask_indices(
                (batch_size, sequence_length),
                mask_prob=self.config.mask_time_prob,
                mask_length=self.config.mask_time_length,
                min_masks=2,
            )
            hidden_states = tf.where(
                tf.cast(mask_time_indices[:, :, tf.newaxis], tf.bool),
                self.masked_spec_embed[tf.newaxis, tf.newaxis, :],
                hidden_states,
            )

        # apply SpecAugment along feature axis
        if self.config.mask_feature_prob > 0:
            mask_feature_indices = _compute_mask_indices(
                (batch_size, hidden_size),
                mask_prob=self.config.mask_feature_prob,
                mask_length=self.config.mask_feature_length,
            )
            hidden_states = tf.where(mask_feature_indices[:, tf.newaxis, :], hidden_states, 0)

        return hidden_states
//  [python]
def call(
        self,
        input_values: tf.Tensor,
        attention_mask: tf.Tensor | None = None,
        token_type_ids: tf.Tensor | None = None,
        position_ids: tf.Tensor | None = None,
        head_mask: tf.Tensor | None = None,
        inputs_embeds: tf.Tensor | None = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        training: bool = False,
        **kwargs: Any,
    ):
        extract_features = self.feature_extractor(t.convert_to_tensor(input_values, dtype=tf.float32), training=training)
        # extract_features = tf.transpose(extract_features, perm=(0, 2, 1))

        if attention_mask is not None:
            # compute real output lengths according to convolution formula
            output_lengths = self._get_feat_extract_output_lengths(tf.reduce_sum(attention_mask, -1))

            attention_mask = tf.sequence_mask(
                output_lengths, maxlen=shape_list(extract_features)[1], dtype=extract_features.dtype
            )

        hidden_states, extract_features = self.feature_projection(extract_features, training=training)

        mask_time_indices = kwargs.get("mask_time_indices", None)
        if training:
            hidden_states = self._mask_hidden_states(hidden_states, mask_time_indices=mask_time_indices)

        encoder_outputs = self.encoder(
            hidden_states,
            attention_mask=attention_mask,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
            training=training,
        )
        hidden_states = encoder_outputs[0]

        if not return_dict:
            return (hidden_states, extract_features) + encoder_outputs[1:]

        return TFWav2Vec2BaseModelOutput(
            last_hidden_state=hidden_states,
            extract_features=extract_features,
            hidden_states=encoder_outputs.hidden_states,
            attentions=encoder_outputs.attentions,
        )
//  [python]
def call(
        self,
        input_values: tf.Tensor,
        attention_mask: tf.Tensor | None = None,
        token_type_ids: tf.Tensor | None = None,
        position_ids: tf.Tensor | None = None,
        head_mask: tf.Tensor | None = None,
        inputs_embeds: tf.Tensor | None = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        training: bool = False,
        **kwargs: Any,
    ):
        extract_features = self.feature_extractor(tf.cast(input_values, tf.float32), training=training)
        # extract_features = tf.transpose(extract_features, perm=(0, 2, 1))

        if attention_mask is not None:
            # compute real output lengths according to convolution formula
            output_lengths = self._get_feat_extract_output_lengths(tf.reduce_sum(attention_mask, -1))

            attention_mask = tf.sequence_mask(
                output_lengths, maxlen=shape_list(extract_features)[1], dtype=extract_features.dtype
            )

        hidden_states, extract_features = self.feature_projection(extract_features, training=training)

        mask_time_indices = kwargs.get("mask_time_indices", None)
        if training:
            hidden_states = self._mask_hidden_states(hidden_states, mask_time_indices=mask_time_indices)

        encoder_outputs = self.encoder(
            hidden_states,
            attention_mask=attention_mask,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
            training=training,
        )
        hidden_states = encoder_outputs[0]

        if not return_dict:
            return (hidden_states, extract_features) + encoder_outputs[1:]

        return TFWav2Vec2BaseModelOutput(
            last_hidden_state=hidden_states,
            extract_features=extract_features,
            hidden_states=encoder_outputs.hidden_states,
            attentions=encoder_outputs.attentions,
        )
//  [python]
def _normalize(self, A, eps=1e-6):
        n = tf.shape(A)[-1]
        A -= tf.linalg.diag(tf.linalg.diag_part(A))
        A_hat = A + t.eye(n, dtype=A.dtype)[tf.newaxis, :, :]
        deg = tf.reduce_sum(A_hat, axis=2)

        D_symm = tf.linalg.diag(1./(eps + tf.math.sqrt(deg)))
        D_asymm = tf.linalg.diag(1./(eps + deg))

        return [A, tf.matmul(D_asymm, A_hat), tf.matmul(tf.matmul(D_symm, A_hat), D_symm)]
//  [python]
def _normalize(self, A, eps=1e-6):
        n = tf.shape(A)[-1]
        A -= tf.linalg.diag(tf.linalg.diag_part(A))
        A_hat = A + tf.cast(tf.eye(n), dtype=A.dtype)[tf.newaxis, :, :]
        deg = tf.reduce_sum(A_hat, axis=2)

        D_symm = tf.linalg.diag(1./(eps + tf.math.sqrt(deg)))
        D_asymm = tf.linalg.diag(1./(eps + deg))

        return [A, tf.matmul(D_asymm, A_hat), tf.matmul(tf.matmul(D_symm, A_hat), D_symm)]
//  [python]
def __init__(self, inner_optimizer, loss_scaling):
        """
        Constructs loss scale optimizer.
        :param inner_optimizer: The optimizer to wrap with this
        functionality.
        :param loss_scaling: Constant loss scaling factor to be applied.
        """
        super().__init__(inner_optimizer, dynamic=False, initial_scale=loss_scaling, dynamic_growth_steps=None)
        self.inner_optimizer.gradient_transformers.insert(
            0, lambda grads_and_vars: [(t.cast(g, dtype=t.float32) / loss_scaling, v) for g, v in grads_and_vars]
        )
//  [python]
def __init__(self, inner_optimizer, loss_scaling):
        """
        Constructs loss scale optimizer.
        :param inner_optimizer: The optimizer to wrap with this
        functionality.
        :param loss_scaling: Constant loss scaling factor to be applied.
        """
        super().__init__(inner_optimizer, dynamic=False, initial_scale=loss_scaling, dynamic_growth_steps=None)
        self.inner_optimizer.gradient_transformers.insert(
            0, lambda grads_and_vars: [(tf.cast(g, dtype=tf.float32) / loss_scaling, v) for g, v in grads_and_vars]
        )
//  [python]
def select_random(x):
  """Selectes a random elements from each row of x."""
  def to_float(x):
    return t.cast(x, t.float32)
  def to_int(x):
    return t.cast(x, t.int64)
  batch_size = tf.shape(x)[0]
  rn = tf.range(batch_size)
  nnz = to_float(tf.count_nonzero(x >= 0, axis=1))
  rnd = tf.random_uniform([batch_size])
  ids = tf.stack([to_int(rn), to_int(nnz * rnd)], axis=1)
  return to_int(tf.gather_nd(x, ids))
//  [python]
def select_random(x):
  """Selectes a random elements from each row of x."""
  def to_float(x):
    return tf.cast(x, tf.float32)
  def to_int(x):
    return tf.cast(x, tf.int64)
  batch_size = tf.shape(x)[0]
  rn = tf.range(batch_size)
  nnz = to_float(tf.count_nonzero(x >= 0, axis=1))
  rnd = tf.random_uniform([batch_size])
  ids = tf.stack([to_int(rn), to_int(nnz * rnd)], axis=1)
  return to_int(tf.gather_nd(x, ids))
//  [python]
def call(
        self,
        input_values: tf.Tensor,
        attention_mask: tf.Tensor | None = None,
        output_attentions: bool | None = None,
        output_hidden_states: bool | None = None,
        return_dict: bool | None = None,
        labels: tf.Tensor | None = None,
        training: bool = False,
    ) -> TFSequenceClassifierOutput | Tuple[tf.Tensor]:
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
        output_hidden_states = True if self.config.use_weighted_layer_sum else output_hidden_states

        outputs = self.wav2vec2(
            input_values,
            attention_mask=attention_mask,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
            training=training,
        )
        if self.config.use_weighted_layer_sum:
            hidden_states = outputs[_HIDDEN_STATES_START_POSITION]
            hidden_states = tf.stack(hidden_states, axis=1)
            norm_weights = tf.nn.softmax(self.layer_weights, axis=-1)
            hidden_states = tf.reduce_sum(hidden_states * tf.reshape(norm_weights, [-1, 1, 1]), axis=1)
        else:
            hidden_states = outputs[0]

        hidden_states = self.projector(hidden_states)
        if attention_mask is None:
            pooled_output = tf.reduce_mean(hidden_states, axis=1)
        else:
            padding_mask = self._get_feature_vector_attention_mask(shape_list(hidden_states)[1], attention_mask)
            padding_mask_float = tf.to_float(padding_mask)
            hidden_states = tf.multiply(hidden_states, tf.expand_dims(padding_mask_float, axis=-1))
            pooled_output = tf.divide(
                tf.reduce_sum(hidden_states, axis=1), tf.expand_dims(tf.reduce_sum(padding_mask_float, axis=1), axis=1)
            )
        logits = self.classifier(pooled_output)
        loss = None
        if labels is not None:
            loss_fn = keras.losses.SparseCategoricalCrossentropy(from_logits=True)
            loss = loss_fn(tf.reshape(labels, [-1]), tf.reshape(logits, [-1, self.config.num_labels]))
        if not return_dict:
            output = (logits,) + outputs[_HIDDEN_STATES_START_POSITION:]
            return ((loss,) + output) if loss is not None else output

        return TFSequenceClassifierOutput(
            loss=loss,
            logits=logits,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
        )
//  [python]
def call(
        self,
        input_values: tf.Tensor,
        attention_mask: tf.Tensor | None = None,
        output_attentions: bool | None = None,
        output_hidden_states: bool | None = None,
        return_dict: bool | None = None,
        labels: tf.Tensor | None = None,
        training: bool = False,
    ) -> TFSequenceClassifierOutput | Tuple[tf.Tensor]:
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
        output_hidden_states = True if self.config.use_weighted_layer_sum else output_hidden_states

        outputs = self.wav2vec2(
            input_values,
            attention_mask=attention_mask,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
            training=training,
        )
        if self.config.use_weighted_layer_sum:
            hidden_states = outputs[_HIDDEN_STATES_START_POSITION]
            hidden_states = tf.stack(hidden_states, axis=1)
            norm_weights = tf.nn.softmax(self.layer_weights, axis=-1)
            hidden_states = tf.reduce_sum(hidden_states * tf.reshape(norm_weights, [-1, 1, 1]), axis=1)
        else:
            hidden_states = outputs[0]

        hidden_states = self.projector(hidden_states)
        if attention_mask is None:
            pooled_output = tf.reduce_mean(hidden_states, axis=1)
        else:
            padding_mask = self._get_feature_vector_attention_mask(shape_list(hidden_states)[1], attention_mask)
            padding_mask_float = tf.cast(padding_mask, hidden_states.dtype)
            hidden_states = tf.multiply(hidden_states, tf.expand_dims(padding_mask_float, axis=-1))
            pooled_output = tf.divide(
                tf.reduce_sum(hidden_states, axis=1), tf.expand_dims(tf.reduce_sum(padding_mask_float, axis=1), axis=1)
            )
        logits = self.classifier(pooled_output)
        loss = None
        if labels is not None:
            loss_fn = keras.losses.SparseCategoricalCrossentropy(from_logits=True)
            loss = loss_fn(tf.reshape(labels, [-1]), tf.reshape(logits, [-1, self.config.num_labels]))
        if not return_dict:
            output = (logits,) + outputs[_HIDDEN_STATES_START_POSITION:]
            return ((loss,) + output) if loss is not None else output

        return TFSequenceClassifierOutput(
            loss=loss,
            logits=logits,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
        )
//  [python]
def _rescale_weights(self, neuron_mask, var, new_weights, axis, scaler):
    non_dead_count = neuron_mask.shape[0] - tf.reduce_sum(neuron_mask)
    non_recycled_norm = tf.reduce_sum(
        (tf.norm(var, axis=axis) * t.cast(1 - neuron_mask, tf.float32))
    ) / t.cast(non_dead_count, tf.float32)

    new_recycled_norm = non_recycled_norm * scaler
    weights_norm_per_neuron = tf.norm(new_weights, axis=axis, keepdims=True)
    new_weights = new_weights / weights_norm_per_neuron
    new_weights = new_recycled_norm * new_weights
    return new_weights
//  [python]
def _rescale_weights(self, neuron_mask, var, new_weights, axis, scaler):
    non_dead_count = neuron_mask.shape[0] - tf.reduce_sum(neuron_mask)
    non_recycled_norm = tf.reduce_sum(
        (tf.norm(var, axis=axis) * tf.cast(1 - neuron_mask, tf.float32))
    ) / tf.cast(non_dead_count, tf.float32)

    new_recycled_norm = non_recycled_norm * scaler
    weights_norm_per_neuron = tf.norm(new_weights, axis=axis, keepdims=True)
    new_weights = new_weights / weights_norm_per_neuron
    new_weights = new_recycled_norm * new_weights
    return new_weights
//  [python]
def hf_compute_loss(self, logits, labels):
        # upsample logits to the images' original size
        # `labels` is of shape (batch_size, height, width)
        label_interp_shape = shape_list(labels)[1:]

        upsampled_logits = tf.image.resize(logits, size=label_interp_shape, method="bilinear")
        # compute weighted loss
        loss_fct = keras.losses.SparseCategoricalCrossentropy(from_logits=True, reduction="none")

        def masked_loss(real, pred):
            unmasked_loss = loss_fct(real, pred)
            mask = tf.to_float(real != self.config.semantic_loss_ignore_index)
            masked_loss = unmasked_loss * mask
            # Reduction strategy in the similar spirit with
            # https://github.com/huggingface/transformers/blob/main/src/transformers/modeling_tf_utils.py#L210
            reduced_masked_loss = tf.reduce_sum(masked_loss) / tf.reduce_sum(mask)
            return tf.reshape(reduced_masked_loss, (1,))

        return masked_loss(labels, upsampled_logits)
//  [python]
def hf_compute_loss(self, logits, labels):
        # upsample logits to the images' original size
        # `labels` is of shape (batch_size, height, width)
        label_interp_shape = shape_list(labels)[1:]

        upsampled_logits = tf.image.resize(logits, size=label_interp_shape, method="bilinear")
        # compute weighted loss
        loss_fct = keras.losses.SparseCategoricalCrossentropy(from_logits=True, reduction="none")

        def masked_loss(real, pred):
            unmasked_loss = loss_fct(real, pred)
            mask = tf.cast(real != self.config.semantic_loss_ignore_index, dtype=unmasked_loss.dtype)
            masked_loss = unmasked_loss * mask
            # Reduction strategy in the similar spirit with
            # https://github.com/huggingface/transformers/blob/main/src/transformers/modeling_tf_utils.py#L210
            reduced_masked_loss = tf.reduce_sum(masked_loss) / tf.reduce_sum(mask)
            return tf.reshape(reduced_masked_loss, (1,))

        return masked_loss(labels, upsampled_logits)
//  [python]
def call(self, encoder_hidden_states: tf.Tensor, training: bool = False) -> tf.Tensor:
        all_hidden_states = ()
        for encoder_hidden_state, mlp in zip(encoder_hidden_states, self.mlps):
            if self.config.reshape_last_stage is False and len(shape_list(encoder_hidden_state)) == 3:
                height = tf.math.sqrt(tf.to_float(shape_list(encoder_hidden_state)[1]))
                height = width = tf.to_int32(height)
                channel_dim = shape_list(encoder_hidden_state)[-1]
                encoder_hidden_state = tf.reshape(encoder_hidden_state, (-1, height, width, channel_dim))

            # unify channel dimension
            encoder_hidden_state = tf.transpose(encoder_hidden_state, perm=[0, 2, 3, 1])
            height, width = shape_list(encoder_hidden_state)[1:3]
            encoder_hidden_state = mlp(encoder_hidden_state)
            channel_dim = shape_list(encoder_hidden_state)[-1]
            encoder_hidden_state = tf.reshape(encoder_hidden_state, (-1, height, width, channel_dim))

            # upsample
            temp_state = tf.transpose(encoder_hidden_states[0], perm=[0, 2, 3, 1])
            upsample_resolution = shape_list(temp_state)[1:-1]
            encoder_hidden_state = tf.image.resize(encoder_hidden_state, size=upsample_resolution, method="bilinear")
            all_hidden_states += (encoder_hidden_state,)

        hidden_states = self.linear_fuse(tf.concat(all_hidden_states[::-1], axis=-1))
        hidden_states = self.batch_norm(hidden_states, training=training)
        hidden_states = self.activation(hidden_states)
        hidden_states = self.dropout(hidden_states, training=training)

        # logits of shape (batch_size, height/4, width/4, num_labels)
        logits = self.classifier(hidden_states)

        return logits
//  [python]
def call(self, encoder_hidden_states: tf.Tensor, training: bool = False) -> tf.Tensor:
        all_hidden_states = ()
        for encoder_hidden_state, mlp in zip(encoder_hidden_states, self.mlps):
            if self.config.reshape_last_stage is False and len(shape_list(encoder_hidden_state)) == 3:
                height = tf.math.sqrt(tf.cast(shape_list(encoder_hidden_state)[1], tf.float32))
                height = width = tf.cast(height, tf.int32)
                channel_dim = shape_list(encoder_hidden_state)[-1]
                encoder_hidden_state = tf.reshape(encoder_hidden_state, (-1, height, width, channel_dim))

            # unify channel dimension
            encoder_hidden_state = tf.transpose(encoder_hidden_state, perm=[0, 2, 3, 1])
            height, width = shape_list(encoder_hidden_state)[1:3]
            encoder_hidden_state = mlp(encoder_hidden_state)
            channel_dim = shape_list(encoder_hidden_state)[-1]
            encoder_hidden_state = tf.reshape(encoder_hidden_state, (-1, height, width, channel_dim))

            # upsample
            temp_state = tf.transpose(encoder_hidden_states[0], perm=[0, 2, 3, 1])
            upsample_resolution = shape_list(temp_state)[1:-1]
            encoder_hidden_state = tf.image.resize(encoder_hidden_state, size=upsample_resolution, method="bilinear")
            all_hidden_states += (encoder_hidden_state,)

        hidden_states = self.linear_fuse(tf.concat(all_hidden_states[::-1], axis=-1))
        hidden_states = self.batch_norm(hidden_states, training=training)
        hidden_states = self.activation(hidden_states)
        hidden_states = self.dropout(hidden_states, training=training)

        # logits of shape (batch_size, height/4, width/4, num_labels)
        logits = self.classifier(hidden_states)

        return logits
//  [python]
def attention(self, inputs_a, inputs_b, attention_size, params, mask=None, return_alphas=False):
        """
        inputs_a = (b, 18, 100)
        inputs_b = (b, 100)
        :param inputs_a:
        :param inputs_b:
        :param attention_size:
        :param time_major:
        :param return_alphas:
        :return:
        """
        if mask is not None:
            mask = tf.cast(self.mask, tf.bool)
        shared = True
        if shared:
            scope_name = 'attn'
        else:
            scope_name = 'attn_'
        with tf.variable_scope(scope_name, reuse=tf.AUTO_REUSE):
            hidden_size = inputs_a.shape[2].value  # D value - hidden size of the RNN layer
            den = False
            x_proj = inputs_a
            y_proj = inputs_b
            # print('x_proj', x_proj.get_shape())
            # print('y_proj', y_proj.get_shape())

            # Trainable parameters
            w_omega = params['w_omega']
            b_omega = params['b_omega']
            # dense_attention_2 = params['dense']
            with tf.variable_scope('v', reuse=tf.AUTO_REUSE):
                # Applying fully connected layer with non-linear activation to each of the B*T timestamps;
                #  the shape of `v` is (B,T,D)*(D,A)=(B,T,A), where A=attention_size

                v = tf.tensordot(x_proj, w_omega, axes=1) + b_omega
                # v  = dense_attention_2(x_proj)

            # For each of the timestamps its vector of size A from `v` is reduced with `u` vector
            vu = tf.tanh(tf.matmul(v, tf.expand_dims(y_proj, -1), name='vu'))  # (B,T) shape (B T A) * (B A 1) = (B T)
            vu = tf.squeeze(vu, -1)
            # print('vu', vu.get_shape())
            # masking
            # mask = None
            if mask is not None:
                vu = tf.where(mask, vu, tf.zeros(tf.shape(vu), dtype=tf.float32))

            alphas = tf.nn.softmax(vu, 1, name='alphas')  # (B,T) shape
            if mask is not None:
                alphas = tf.where(mask, alphas, tf.zeros(tf.shape(alphas), dtype=tf.float32))
                a = tf.reduce_sum(tf.expand_dims(alphas, -1), axis=1)
                condition = tf.equal(a, 0.0)
                case_true = tf.ones(tf.shape(a), tf.float32)
                a_m = tf.where(condition, case_true, a)
                alphas = tf.divide(alphas, a_m)

            # print('alphas', alphas.get_shape())

            # Output of (Bi-)RNN is reduced with attention vector; the result has (B,D) shape
            output = tf.matmul(tf.transpose(inputs_a, [0, 2, 1]), tf.expand_dims(alphas, -1))
            output = tf.squeeze(output, -1)
            # print('r', output.get_shape())
            # output = tf.reduce_sum(r, 1)

            if not return_alphas:
                return tf.expand_dims(output, 1)
            else:
                return tf.expand_dims(output, 1), alphas
//  [python]
def hard_sigmoid(features):
  """Computes the hard sigmoid activation function.

  Args:
    features: A `Tensor` representing preactivation values.

  Returns:
    The activation value.
  """
  features = tf.convert_to_tensor(features)
  return tf.nn.relu6(features + tf.to_float(3., features.dtype)) * 0.16667
//  [python]
def hard_sigmoid(features):
  """Computes the hard sigmoid activation function.

  Args:
    features: A `Tensor` representing preactivation values.

  Returns:
    The activation value.
  """
  features = tf.convert_to_tensor(features)
  return tf.nn.relu6(features + tf.cast(3., features.dtype)) * 0.16667
//  [python]
def get_optimizer(
    optimizer_name,
    num_replicas,
    learning_rate_schedule,
    use_outlining=False,
    loss_scaling=None,
    beta1=0.9,
    beta2=0.999,
    epsilon=1e-6,
    weight_decay_rate=0.0,
):
    """Constructs and returns a Keras optimizer
    Args:
        optimizer_name: A string representing the name of the
            required optimizer
        num_replicas: An integer representing the number of replicas. This
            will be used to scale down the gradients.
        learning_rate_schedule: A float or a Keras learning rate schedule
            object.
        use_outlining: A boolean, if true, the optimizer update will be
            outlined.
        loss_scaling: A float representing the fixed loss scaling. If None,
            the loss will not be scaled.
        beta1: A `float` value or a constant `float` tensor.
            The exponential decay rate for the 1st moment estimates.
        beta2: A `float` value or a constant `float` tensor.
            The exponential decay rate for the 2nd moment estimates.
        epsilon: A small constant for numerical stability.
        weight_decay_rate: A `float` value to decay the variables by in the
            gradient update step.
    """

    def rescale_gradients_down(grads_and_vars):
        return [(t.cast(g, dtype=t.float32) / num_replicas, v) for g, v in grads_and_vars]

    optimizer_kwargs = {
        "learning_rate": learning_rate_schedule,
        "beta_1": beta1,
        "beta_2": beta2,
        "epsilon": epsilon,
        "m_dtype": t.float32,
        "v_dtype": t.float32,
        "gradient_transformers": [rescale_gradients_down],
        "outline_apply_gradients": use_outlining,
        "outline_apply_gradients_kwargs": {"unique_sharding": True} if use_outlining else None,
    }

    if optimizer_name == "AdamW":
        AdamWIpuOptimizer = extend_with_decoupled_weight_decay(AdamIpuOptimizer)
        optimizer = AdamWIpuOptimizer(
            weight_decay_rate,
            optimizer_compute_precisions=(t.float32,),
            **optimizer_kwargs,
        )
    elif optimizer_name == "LAMB":
        optimizer = LAMBIpuOptimizer(
            weight_decay_rate=weight_decay_rate,
            exclude_from_layer_adaptation=["bias", "beta", "gamma"],
            exclude_from_weight_decay=["bias", "beta"],
            optimizer_compute_precisions=(t.float32, t.float16),
            **optimizer_kwargs,
        )
    else:
        raise ValueError(f"Unrecognised optimizer name: `{optimizer_name}`." f" Choose one of {ALLOWED_OPTIMIZERS}")

    if loss_scaling is not None:
        optimizer = StaticLossScaleOptimizer(optimizer, loss_scaling)

    return optimizer
//  [python]
def get_optimizer(
    optimizer_name,
    num_replicas,
    learning_rate_schedule,
    use_outlining=False,
    loss_scaling=None,
    beta1=0.9,
    beta2=0.999,
    epsilon=1e-6,
    weight_decay_rate=0.0,
):
    """Constructs and returns a Keras optimizer
    Args:
        optimizer_name: A string representing the name of the
            required optimizer
        num_replicas: An integer representing the number of replicas. This
            will be used to scale down the gradients.
        learning_rate_schedule: A float or a Keras learning rate schedule
            object.
        use_outlining: A boolean, if true, the optimizer update will be
            outlined.
        loss_scaling: A float representing the fixed loss scaling. If None,
            the loss will not be scaled.
        beta1: A `float` value or a constant `float` tensor.
            The exponential decay rate for the 1st moment estimates.
        beta2: A `float` value or a constant `float` tensor.
            The exponential decay rate for the 2nd moment estimates.
        epsilon: A small constant for numerical stability.
        weight_decay_rate: A `float` value to decay the variables by in the
            gradient update step.
    """

    def rescale_gradients_down(grads_and_vars):
        return [(tf.cast(g, dtype=tf.float32) / num_replicas, v) for g, v in grads_and_vars]

    optimizer_kwargs = {
        "learning_rate": learning_rate_schedule,
        "beta_1": beta1,
        "beta_2": beta2,
        "epsilon": epsilon,
        "m_dtype": tf.float32,
        "v_dtype": tf.float32,
        "gradient_transformers": [rescale_gradients_down],
        "outline_apply_gradients": use_outlining,
        "outline_apply_gradients_kwargs": {"unique_sharding": True} if use_outlining else None,
    }

    if optimizer_name == "AdamW":
        AdamWIpuOptimizer = extend_with_decoupled_weight_decay(AdamIpuOptimizer)
        optimizer = AdamWIpuOptimizer(
            weight_decay_rate,
            optimizer_compute_precisions=(tf.float32,),
            **optimizer_kwargs,
        )
    elif optimizer_name == "LAMB":
        optimizer = LAMBIpuOptimizer(
            weight_decay_rate=weight_decay_rate,
            exclude_from_layer_adaptation=["bias", "beta", "gamma"],
            exclude_from_weight_decay=["bias", "beta"],
            optimizer_compute_precisions=(tf.float32, tf.float16),
            **optimizer_kwargs,
        )
    else:
        raise ValueError(f"Unrecognised optimizer name: `{optimizer_name}`." f" Choose one of {ALLOWED_OPTIMIZERS}")

    if loss_scaling is not None:
        optimizer = StaticLossScaleOptimizer(optimizer, loss_scaling)

    return optimizer
//  [python]
def _generate_synthetic_data(self, ctx: tf.distribute.InputContext,
                               batch_size: int) -> tf.data.Dataset:
    """Creates synthetic data based on the parameter batch size.

    Args:
      ctx: Input Context
      batch_size: per replica batch size.

    Returns:
      The synthetic dataset.
    """
    params = self._params
    num_dense = self._num_dense_features
    num_replicas = ctx.num_replicas_in_sync if ctx else 1

    if params.is_training:
      dataset_size = 50 * batch_size * num_replicas
    else:
      dataset_size = 50 * batch_size * num_replicas
    dense_tensor = tf.random.uniform(
        shape=(dataset_size, num_dense), maxval=1.0, dtype=tf.float32
    )

    sparse_tensors = []
    for vocab_size, multi_hot_size in zip(
        self._vocab_sizes, self._multi_hot_sizes
    ):
      sparse_tensors.append(
          tf.sparse.from_dense(
              tf.random.uniform(
                  shape=(dataset_size, multi_hot_size),
                  maxval=int(vocab_size),
                  dtype=tf.int32,
              )
          )
      )

    sparse_tensor_elements = {
        str(i): sparse_tensors[i] for i in range(len(sparse_tensors))
    }

    # the mean is in [0, 1] interval.
    dense_tensor_mean = tf.math.reduce_mean(dense_tensor, axis=1)

    # the label is in [0, 1] interval.
    label_tensor = (dense_tensor_mean)
    # Using the threshold 0.5 to convert to 0/1 labels.
    label_tensor = tf.to_int32(label_tensor + 0.5)

    input_elem = {'dense_features': dense_tensor,
                  'sparse_features': sparse_tensor_elements}, label_tensor

    dataset = tf.data.Dataset.from_tensor_slices(input_elem)
    dataset = dataset.cache()
    if params.is_training:
      dataset = dataset.repeat()

    return dataset.batch(batch_size, drop_remainder=True)
//  [python]
def _generate_synthetic_data(self, ctx: tf.distribute.InputContext,
                               batch_size: int) -> tf.data.Dataset:
    """Creates synthetic data based on the parameter batch size.

    Args:
      ctx: Input Context
      batch_size: per replica batch size.

    Returns:
      The synthetic dataset.
    """
    params = self._params
    num_dense = self._num_dense_features
    num_replicas = ctx.num_replicas_in_sync if ctx else 1

    if params.is_training:
      dataset_size = 50 * batch_size * num_replicas
    else:
      dataset_size = 50 * batch_size * num_replicas
    dense_tensor = tf.random.uniform(
        shape=(dataset_size, num_dense), maxval=1.0, dtype=tf.float32
    )

    sparse_tensors = []
    for vocab_size, multi_hot_size in zip(
        self._vocab_sizes, self._multi_hot_sizes
    ):
      sparse_tensors.append(
          tf.sparse.from_dense(
              tf.random.uniform(
                  shape=(dataset_size, multi_hot_size),
                  maxval=int(vocab_size),
                  dtype=tf.int32,
              )
          )
      )

    sparse_tensor_elements = {
        str(i): sparse_tensors[i] for i in range(len(sparse_tensors))
    }

    # the mean is in [0, 1] interval.
    dense_tensor_mean = tf.math.reduce_mean(dense_tensor, axis=1)

    # the label is in [0, 1] interval.
    label_tensor = (dense_tensor_mean)
    # Using the threshold 0.5 to convert to 0/1 labels.
    label_tensor = tf.cast(label_tensor + 0.5, tf.int32)

    input_elem = {'dense_features': dense_tensor,
                  'sparse_features': sparse_tensor_elements}, label_tensor

    dataset = tf.data.Dataset.from_tensor_slices(input_elem)
    dataset = dataset.cache()
    if params.is_training:
      dataset = dataset.repeat()

    return dataset.batch(batch_size, drop_remainder=True)
//  [python]
def scale(boxlist, y_scale, x_scale, scope=None):
  """scale box coordinates in x and y dimensions.

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

  Returns:
    boxlist: BoxList holding N boxes
  """
  with tf.name_scope(scope, 'Scale'):
    y_scale = tf.to_float(y_scale)
    x_scale = tf.to_float(x_scale)
    y_min, x_min, y_max, x_max = tf.split(
        value=boxlist.get(), num_or_size_splits=4, axis=1)
    y_min = y_scale * y_min
    y_max = y_scale * y_max
    x_min = x_scale * x_min
    x_max = x_scale * x_max
    scaled_boxlist = box_list.BoxList(
        tf.concat([y_min, x_min, y_max, x_max], 1))
    return _copy_extra_fields(scaled_boxlist, boxlist)
//  [python]
def scale(boxlist, y_scale, x_scale, scope=None):
  """scale box coordinates in x and y dimensions.

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

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

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

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

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

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

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

  Returns:
    image_and_boxes: an image tensor with shape [height, width, 3]
  """
  with tf.name_scope(scope, 'VisualizeBoxesInImage'):
    if not normalized:
      height, width, _ = tf.unstack(tf.shape(image))
      boxlist = scale(boxlist,
                      1.0 / tf.cast(height, tf.float32),
                      1.0 / tf.cast(width, tf.float32))
    corners = tf.expand_dims(boxlist.get(), 0)
    image = tf.expand_dims(image, 0)
    return tf.squeeze(tf.image.draw_bounding_boxes(image, corners), [0])
//  [python]
def _data_with_indices_dtype(
      cls, data: gp.Data, dtype: tf.dtypes.DType
  ) -> gp.Data:
    data = data.copy()
    sizes = data.pop(_GraphPieceWithFeatures._DATAKEY_SIZES)
    data = gp.data_with_indices_dtype(data, dtype)
    data[_GraphPieceWithFeatures._DATAKEY_SIZES] = t.cast(sizes, dtype)
    return data
//  [python]
def _data_with_indices_dtype(
      cls, data: gp.Data, dtype: tf.dtypes.DType
  ) -> gp.Data:
    data = data.copy()
    sizes = data.pop(_GraphPieceWithFeatures._DATAKEY_SIZES)
    data = gp.data_with_indices_dtype(data, dtype)
    data[_GraphPieceWithFeatures._DATAKEY_SIZES] = tf.cast(sizes, dtype)
    return data
//  [python]
def filter_greater_than(boxlist, thresh, scope=None):
  """Filter to keep only boxes with score exceeding a given threshold.

  This op keeps the collection of boxes whose corresponding scores are
  greater than the input threshold.

  TODO(jonathanhuang): Change function name to filter_scores_greater_than

  Args:
    boxlist: BoxList holding N boxes.  Must contain a 'scores' field
      representing detection scores.
    thresh: scalar threshold
    scope: name scope.

  Returns:
    a BoxList holding M boxes where M <= N

  Raises:
    ValueError: if boxlist not a BoxList object or if it does not
      have a scores field
  """
  with tf.name_scope(scope, 'FilterGreaterThan'):
    if not isinstance(boxlist, box_list.BoxList):
      raise ValueError('boxlist must be a BoxList')
    if not boxlist.has_field('scores'):
      raise ValueError('input boxlist must have \'scores\' field')
    scores = boxlist.get_field('scores')
    if len(scores.shape.as_list()) > 2:
      raise ValueError('Scores should have rank 1 or 2')
    if len(scores.shape.as_list()) == 2 and scores.shape.as_list()[1] != 1:
      raise ValueError('Scores should have rank 1 or have shape '
                       'consistent with [None, 1]')
    high_score_indices = tf.reshape(
        tf.where(tf.greater(scores, thresh)),
        [-1])
    return gather(boxlist, tf.to_int32(high_score_indices))
//  [python]
def filter_greater_than(boxlist, thresh, scope=None):
  """Filter to keep only boxes with score exceeding a given threshold.

  This op keeps the collection of boxes whose corresponding scores are
  greater than the input threshold.

  TODO(jonathanhuang): Change function name to filter_scores_greater_than

  Args:
    boxlist: BoxList holding N boxes.  Must contain a 'scores' field
      representing detection scores.
    thresh: scalar threshold
    scope: name scope.

  Returns:
    a BoxList holding M boxes where M <= N

  Raises:
    ValueError: if boxlist not a BoxList object or if it does not
      have a scores field
  """
  with tf.name_scope(scope, 'FilterGreaterThan'):
    if not isinstance(boxlist, box_list.BoxList):
      raise ValueError('boxlist must be a BoxList')
    if not boxlist.has_field('scores'):
      raise ValueError('input boxlist must have \'scores\' field')
    scores = boxlist.get_field('scores')
    if len(scores.shape.as_list()) > 2:
      raise ValueError('Scores should have rank 1 or 2')
    if len(scores.shape.as_list()) == 2 and scores.shape.as_list()[1] != 1:
      raise ValueError('Scores should have rank 1 or have shape '
                       'consistent with [None, 1]')
    high_score_indices = tf.cast(tf.reshape(
        tf.where(tf.greater(scores, thresh)),
        [-1]), tf.int32)
    return gather(boxlist, high_score_indices)
//  [python]
def to_absolute_coordinates(boxlist,
                            height,
                            width,
                            check_range=True,
                            maximum_normalized_coordinate=1.1,
                            scope=None):
  """Converts normalized box coordinates to absolute pixel coordinates.

  This function raises an assertion failed error when the maximum box coordinate
  value is larger than maximum_normalized_coordinate (in which case coordinates
  are already absolute).

  Args:
    boxlist: BoxList with coordinates in range [0, 1].
    height: Maximum value for height of absolute box coordinates.
    width: Maximum value for width of absolute box coordinates.
    check_range: If True, checks if the coordinates are normalized or not.
    maximum_normalized_coordinate: Maximum coordinate value to be considered
      as normalized, default to 1.1.
    scope: name scope.

  Returns:
    boxlist with absolute coordinates in terms of the image size.

  """
  with tf.name_scope(scope, 'ToAbsoluteCoordinates'):
    height = tf.cast(height, tf.float32)
    width = tf.cast(width, tf.float32)

    # Ensure range of input boxes is correct.
    if check_range:
      box_maximum = tf.reduce_max(boxlist.get())
      max_assert = tf.Assert(
          tf.greater_equal(maximum_normalized_coordinate, box_maximum),
          ['maximum box coordinate value is larger '
           'than %f: ' % maximum_normalized_coordinate, box_maximum])
      with tf.control_dependencies([max_assert]):
        width = tf.identity(width)

    return scale(boxlist, height, width)
//  [python]
def to_normalized_coordinates(boxlist, height, width,
                              check_range=True, scope=None):
  """Converts absolute box coordinates to normalized coordinates in [0, 1].

  Usually one uses the dynamic shape of the image or conv-layer tensor:
    boxlist = box_list_ops.to_normalized_coordinates(boxlist,
                                                     tf.shape(images)[1],
                                                     tf.shape(images)[2]),

  This function raises an assertion failed error at graph execution time when
  the maximum coordinate is smaller than 1.01 (which means that coordinates are
  already normalized). The value 1.01 is to deal with small rounding errors.

  Args:
    boxlist: BoxList with coordinates in terms of pixel-locations.
    height: Maximum value for height of absolute box coordinates.
    width: Maximum value for width of absolute box coordinates.
    check_range: If True, checks if the coordinates are normalized or not.
    scope: name scope.

  Returns:
    boxlist with normalized coordinates in [0, 1].
  """
  with tf.name_scope(scope, 'ToNormalizedCoordinates'):
    height = tf.cast(height, tf.float32)
    width = tf.cast(width, tf.float32)

    if check_range:
      max_val = tf.reduce_max(boxlist.get())
      max_assert = tf.Assert(tf.greater(max_val, 1.01),
                             ['max value is lower than 1.01: ', max_val])
      with tf.control_dependencies([max_assert]):
        width = tf.identity(width)

    return scale(boxlist, 1 / height, 1 / width)
//  [python]
def __init__(self):

        # used to feed training/testing or evaluation iterator as string representation
        self.iterator_handle = tf.placeholder(tf.string, shape=[],name="data_iterator_handle")
        # requires info to data dimensions
        #self.output_shapes = tuple(output_shapes)
        #self.output_types = tuple(output_types)

        # changes behavior of batch normalization
        self.is_train = tf.placeholder_with_default(tf.convert_to_tensor(True, dtype=tf.bool), shape=(), name="is_train")

        def parse(var,dtype):
            if type(var)==dtype: return var
            else: return eval(var)

        # see FLAGS for documentation
        self.kernelsize = parse(FLAGS.kernel,tuple)
        self.num_classes = parse(FLAGS.num_classes,int)
//  [python]
def __init__(self):

        # used to feed training/testing or evaluation iterator as string representation
        self.iterator_handle = tf.placeholder(tf.string, shape=[],name="data_iterator_handle")
        # requires info to data dimensions
        #self.output_shapes = tuple(output_shapes)
        #self.output_types = tuple(output_types)

        # changes behavior of batch normalization
        self.is_train = tf.placeholder_with_default(tf.cast(True, tf.bool), shape=(), name="is_train")

        def parse(var,dtype):
            if type(var)==dtype: return var
            else: return eval(var)

        # see FLAGS for documentation
        self.kernelsize = parse(FLAGS.kernel,tuple)
        self.num_classes = parse(FLAGS.num_classes,int)
//  [python]
def multinomial_sampling(self, pr):
        """Multinomial Sampling of ratings

        Basic mechanics:
        For r classes, we sample r binomial distributions using the rejection method. This is possible
        since each class is statistically independent from the other. Note that this is the same method
        used in numpy's random.multinomial() function.

        1) extract a size r array of random numbers from a uniform distribution (g). As pr is normalized,
        we need to normalize g as well.

        2) For each user and item, compare pr with the reference distribution. Note that the latter needs
        to be the same for ALL the user/item pairs in the dataset, as by assumptions they are sampled
        from a common distribution.

        Args:
            pr (tf.Tensor, float32): A distributions of shape (m, n, r), where m is the number of examples, n the number
                 of features and r the number of classes. pr needs to be normalized, i.e. sum_k p(k) = 1 for all m, at fixed n.
            f (tf.Tensor, float32): Normalized, uniform probability used for comparison.

        Returns:
            tf.Tensor: An (m,n) float32 tensor of sampled rankings from 1 to r.
        """
        g = np.random.uniform(size=pr.shape[2])  # sample from a uniform distribution
        f = tf.convert_to_tensor(
            value=g / g.sum(), dtype=tf.float32
        )  # normalize and convert to tensor

        samp = tf.nn.relu(tf.sign(pr - f))  # apply rejection method

        # get integer index of the rating to be sampled
        v_argmax = t.cast(t.argmax(input=samp, axis=2), "int32")

        # lookup the rating using integer index
        v_samp = t.cast(self.ratings_lookup_table.lookup(v_argmax), "float32")

        return v_samp