jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def train (optimizer, model_static, model_fn,
           train_generator=bilist_generator(([],[])),
           valid_generator=bilist_generator(([],[])),
           epochs=100,
           callbacks=[],
           extraVectors=[]):
    '''
    Train the given model.

    train_generator: training data

    valid_generator: validation data

    epochs: number of epochs

    callbacks: list of callbacks.
      Each callback receives an epoch entry (see below). If it returns False then the training is aborted.

    extraVectors: list of extra vectors to pass to session.run when training.

    modelPrefix: in case of a multitask/multimodel, give the prefix of the model to use.

    This function returns a list of epoch entries. Each entry is a dictionary with:
     - "epoch": current epoch
     - "val" and "train": dictionaries with
        - "loss", "accuracy", "error_rate", time", "start_time", "end_time"
    '''
    batch_size = model_static["batch_size"]
    train_vars = model_static["parameters"]
    placeholders_info = model_fn["placeholders"]
    stats = []
    def halfEpoch(isTraining):
        totalAccur = 0
        totalLoss = 0
        n = 0
        print ("Training" if isTraining else "Validation", end="")
        start_time = time()
        for inputs in train_generator(batch_size) if isTraining else valid_generator(batch_size):
            cast_inputs = dict((k,tf.to_float(inputs[k])) for k in placeholders_info)
            # the above forces inputs to be tensors. (It's convenient to pass just lists here)
            print(".",end="")
            sys.stdout.flush()
            with tf.GradientTape() as tape:
                results = model_fn["function"](tf.constant(isTraining, shape=[]), **{**(model_static["paramsdict"]), **cast_inputs})
                loss = results["loss"]
                accur = results["accuracy"]
                if isTraining:
                   grads = tape.gradient(loss, train_vars)
                   optimizer.apply_gradients(zip(grads, train_vars))
                n+=1
                totalLoss += loss
                totalAccur += accur
        end_time = time()
        totalAccur = totalAccur.numpy()
        totalLoss = totalLoss.numpy()
        if n > 0:
            avgLoss = totalLoss / float(n)
            avgAccur = totalAccur / float(n)
            print(".")
            print ("Time=%.1f" % (end_time - start_time), "loss=%g" % avgLoss, "accuracy=%.3f" % avgAccur)
            return {"loss":avgLoss,"accuracy":avgAccur,"time":(end_time - start_time),"error_rate":1-avgAccur,"start_time":start_time}
        else:
            print ("No data")
            return {"loss":0,"accur":0,"time":0,"error_rate":0,"start_time":0}

    for e in range(epochs):
        print ("Epoch {0}/{1}".format(e, epochs))
        tr = halfEpoch(True)
        va = halfEpoch(False)
        epoch_stats = {"train":tr, "val":va, "epoch":e}
        stats.append(epoch_stats)
        if any(c(epoch_stats) for c in callbacks):
            break
    return stats
//  [python]
def train (optimizer, model_static, model_fn,
           train_generator=bilist_generator(([],[])),
           valid_generator=bilist_generator(([],[])),
           epochs=100,
           callbacks=[],
           extraVectors=[]):
    '''
    Train the given model.

    train_generator: training data

    valid_generator: validation data

    epochs: number of epochs

    callbacks: list of callbacks.
      Each callback receives an epoch entry (see below). If it returns False then the training is aborted.

    extraVectors: list of extra vectors to pass to session.run when training.

    modelPrefix: in case of a multitask/multimodel, give the prefix of the model to use.

    This function returns a list of epoch entries. Each entry is a dictionary with:
     - "epoch": current epoch
     - "val" and "train": dictionaries with
        - "loss", "accuracy", "error_rate", time", "start_time", "end_time"
    '''
    batch_size = model_static["batch_size"]
    train_vars = model_static["parameters"]
    placeholders_info = model_fn["placeholders"]
    stats = []
    def halfEpoch(isTraining):
        totalAccur = 0
        totalLoss = 0
        n = 0
        print ("Training" if isTraining else "Validation", end="")
        start_time = time()
        for inputs in train_generator(batch_size) if isTraining else valid_generator(batch_size):
            cast_inputs = dict((k,tf.cast(inputs[k], placeholders_info[k]["dtype"])) for k in placeholders_info)
            # the above forces inputs to be tensors. (It's convenient to pass just lists here)
            print(".",end="")
            sys.stdout.flush()
            with tf.GradientTape() as tape:
                results = model_fn["function"](tf.constant(isTraining, shape=[]), **{**(model_static["paramsdict"]), **cast_inputs})
                loss = results["loss"]
                accur = results["accuracy"]
                if isTraining:
                   grads = tape.gradient(loss, train_vars)
                   optimizer.apply_gradients(zip(grads, train_vars))
                n+=1
                totalLoss += loss
                totalAccur += accur
        end_time = time()
        totalAccur = totalAccur.numpy()
        totalLoss = totalLoss.numpy()
        if n > 0:
            avgLoss = totalLoss / float(n)
            avgAccur = totalAccur / float(n)
            print(".")
            print ("Time=%.1f" % (end_time - start_time), "loss=%g" % avgLoss, "accuracy=%.3f" % avgAccur)
            return {"loss":avgLoss,"accuracy":avgAccur,"time":(end_time - start_time),"error_rate":1-avgAccur,"start_time":start_time}
        else:
            print ("No data")
            return {"loss":0,"accur":0,"time":0,"error_rate":0,"start_time":0}

    for e in range(epochs):
        print ("Epoch {0}/{1}".format(e, epochs))
        tr = halfEpoch(True)
        va = halfEpoch(False)
        epoch_stats = {"train":tr, "val":va, "epoch":e}
        stats.append(epoch_stats)
        if any(c(epoch_stats) for c in callbacks):
            break
    return stats
//  [python]
def _get_actions_from_reward_layer(
      self, encoded_observation: types.Float, mask: Optional[types.Tensor]
  ) -> Tuple[types.Int, types.Float, types.Float]:
    # Get the predicted expected reward.
    est_mean_reward = tf.reshape(
        self._reward_layer(encoded_observation), shape=[-1, self._num_actions]
    )
    if mask is None:
      greedy_actions = tf.argmax(est_mean_reward, axis=-1, output_type=tf.int32)
    else:
      greedy_actions = policy_utilities.masked_argmax(
          est_mean_reward, mask, output_type=tf.int32
      )

    # Add epsilon greedy on top, if needed.
    if self._epsilon_greedy:
      batch_size = (
          tf.compat.dimension_value(encoded_observation.shape[0])
          or tf.shape(encoded_observation)[0]
      )
      if mask is None:
        random_actions = tf.random.uniform(
            [batch_size], maxval=self._num_actions, dtype=tf.int32
        )
      else:
        zero_logits = tf.to_float(tf.zeros_like(mask))
        masked_categorical = masked.MaskedCategorical(
            zero_logits, mask, dtype=tf.int32
        )
        random_actions = masked_categorical.sample()

      rng = tf.random.uniform([batch_size], maxval=1.0)
      cond = tf.greater(rng, self._epsilon_greedy)
      chosen_actions = tf.compat.v1.where(cond, greedy_actions, random_actions)
    else:
      chosen_actions = greedy_actions

    return chosen_actions, est_mean_reward, est_mean_reward
//  [python]
def _get_actions_from_reward_layer(
      self, encoded_observation: types.Float, mask: Optional[types.Tensor]
  ) -> Tuple[types.Int, types.Float, types.Float]:
    # Get the predicted expected reward.
    est_mean_reward = tf.reshape(
        self._reward_layer(encoded_observation), shape=[-1, self._num_actions]
    )
    if mask is None:
      greedy_actions = tf.argmax(est_mean_reward, axis=-1, output_type=tf.int32)
    else:
      greedy_actions = policy_utilities.masked_argmax(
          est_mean_reward, mask, output_type=tf.int32
      )

    # Add epsilon greedy on top, if needed.
    if self._epsilon_greedy:
      batch_size = (
          tf.compat.dimension_value(encoded_observation.shape[0])
          or tf.shape(encoded_observation)[0]
      )
      if mask is None:
        random_actions = tf.random.uniform(
            [batch_size], maxval=self._num_actions, dtype=tf.int32
        )
      else:
        zero_logits = tf.cast(tf.zeros_like(mask), tf.float32)
        masked_categorical = masked.MaskedCategorical(
            zero_logits, mask, dtype=tf.int32
        )
        random_actions = masked_categorical.sample()

      rng = tf.random.uniform([batch_size], maxval=1.0)
      cond = tf.greater(rng, self._epsilon_greedy)
      chosen_actions = tf.compat.v1.where(cond, greedy_actions, random_actions)
    else:
      chosen_actions = greedy_actions

    return chosen_actions, est_mean_reward, est_mean_reward
//  [python]
def _postprocess_conv2d_output(inputs: tf.Tensor, data_format: str | None) -> tf.Tensor:
        """Transpose and cast the output from conv2d if needed.

        Parameters
        ----------
        inputs: :class:`tf.Tensor`
            The input that requires transposing and casting
        data_format: str
            `"channels_last"` or `"channels_first"`

        Returns
        -------
        :class:`tf.Tensor`
            The transposed and cast input tensor
        """

        if data_format == "channels_first":
            inputs = tf.transpose(inputs, (0, 3, 1, 2))

        if K.floatx() == "float64":
            inputs = tf.cast(inputs, "float64")
        return inputs
//  [python]
def _get_actions_from_linucb(
      self, encoded_observation: types.Float, mask: Optional[types.Tensor]
  ) -> Tuple[types.Int, types.Float, types.Float]:
    encoded_observation = t.cast(encoded_observation, dtype=self._dtype)

    p_values = []
    est_rewards = []
    for k in range(self._num_actions):
      encoded_observation_for_arm = self._get_encoded_observation_for_arm(
          encoded_observation, k
      )
      model_index = policy_utilities.get_model_index(
          k, self._accepts_per_arm_features
      )
      a_inv_x = linalg.conjugate_gradient(
          self._cov_matrix[model_index]
          + tf.eye(self._encoding_dim, dtype=self._dtype),
          tf.linalg.matrix_transpose(encoded_observation_for_arm),
      )
      mean_reward_est = tf.einsum(
          'j,jk->k', self._data_vector[model_index], a_inv_x
      )
      est_rewards.append(mean_reward_est)

      ci = tf.reshape(
          tf.linalg.tensor_diag_part(
              tf.matmul(encoded_observation_for_arm, a_inv_x)
          ),
          [-1, 1],
      )
      p_values.append(
          tf.reshape(mean_reward_est, [-1, 1]) + self._alpha * tf.sqrt(ci)
      )

    stacked_p_values = tf.squeeze(tf.stack(p_values, axis=-1), axis=[1])
    if mask is None:
      chosen_actions = tf.argmax(
          stacked_p_values, axis=-1, output_type=tf.int32
      )
    else:
      chosen_actions = policy_utilities.masked_argmax(
          stacked_p_values, mask, output_type=tf.int32
      )

    est_mean_reward = tf.cast(tf.stack(est_rewards, axis=-1), tf.float32)
    return (
        chosen_actions,
        est_mean_reward,
        tf.cast(stacked_p_values, tf.float32),
    )
//  [python]
def _get_actions_from_linucb(
      self, encoded_observation: types.Float, mask: Optional[types.Tensor]
  ) -> Tuple[types.Int, types.Float, types.Float]:
    encoded_observation = tf.cast(encoded_observation, dtype=self._dtype)

    p_values = []
    est_rewards = []
    for k in range(self._num_actions):
      encoded_observation_for_arm = self._get_encoded_observation_for_arm(
          encoded_observation, k
      )
      model_index = policy_utilities.get_model_index(
          k, self._accepts_per_arm_features
      )
      a_inv_x = linalg.conjugate_gradient(
          self._cov_matrix[model_index]
          + tf.eye(self._encoding_dim, dtype=self._dtype),
          tf.linalg.matrix_transpose(encoded_observation_for_arm),
      )
      mean_reward_est = tf.einsum(
          'j,jk->k', self._data_vector[model_index], a_inv_x
      )
      est_rewards.append(mean_reward_est)

      ci = tf.reshape(
          tf.linalg.tensor_diag_part(
              tf.matmul(encoded_observation_for_arm, a_inv_x)
          ),
          [-1, 1],
      )
      p_values.append(
          tf.reshape(mean_reward_est, [-1, 1]) + self._alpha * tf.sqrt(ci)
      )

    stacked_p_values = tf.squeeze(tf.stack(p_values, axis=-1), axis=[1])
    if mask is None:
      chosen_actions = tf.argmax(
          stacked_p_values, axis=-1, output_type=tf.int32
      )
    else:
      chosen_actions = policy_utilities.masked_argmax(
          stacked_p_values, mask, output_type=tf.int32
      )

    est_mean_reward = tf.cast(tf.stack(est_rewards, axis=-1), tf.float32)
    return (
        chosen_actions,
        est_mean_reward,
        tf.cast(stacked_p_values, tf.float32),
    )
//  [python]
def _action(self, time_step, policy_state, seed):  # pytype: disable=signature-mismatch  # overriding-parameter-count-checks
    observation = time_step.observation
    if self.observation_and_action_constraint_splitter is not None:
      observation, _ = self.observation_and_action_constraint_splitter(
          observation
      )
    mask = constraints.construct_mask_from_multiple_sources(
        time_step.observation,
        self._observation_and_action_constraint_splitter,
        (),
        self._num_actions,
    )
    # Pass the observations through the encoding network.
    encoded_observation, _ = self._encoding_network(observation)
    encoded_observation = t.cast(encoded_observation, dtype=self._dtype)

    if tf.distribute.has_strategy():
      if self._distributed_use_reward_layer:
        chosen_actions, est_mean_rewards, est_rewards_optimistic = (
            self._get_actions_from_reward_layer(encoded_observation, mask)
        )
      else:
        chosen_actions, est_mean_rewards, est_rewards_optimistic = (
            self._get_actions_from_linucb(encoded_observation, mask)
        )
    else:
      chosen_actions, est_mean_rewards, est_rewards_optimistic = tf.cond(
          self._actions_from_reward_layer,
          # pylint: disable=g-long-lambda
          lambda: self._get_actions_from_reward_layer(
              encoded_observation, mask
          ),
          lambda: self._get_actions_from_linucb(encoded_observation, mask),
      )

    arm_observations = ()
    if self._accepts_per_arm_features:
      arm_observations = observation[bandit_spec_utils.PER_ARM_FEATURE_KEY]
    policy_info = policy_utilities.populate_policy_info(
        arm_observations,
        chosen_actions,
        est_rewards_optimistic,
        est_mean_rewards,
        self._emit_policy_info,
        self._accepts_per_arm_features,
    )
    return policy_step.PolicyStep(chosen_actions, policy_state, policy_info)
//  [python]
def _action(self, time_step, policy_state, seed):  # pytype: disable=signature-mismatch  # overriding-parameter-count-checks
    observation = time_step.observation
    if self.observation_and_action_constraint_splitter is not None:
      observation, _ = self.observation_and_action_constraint_splitter(
          observation
      )
    mask = constraints.construct_mask_from_multiple_sources(
        time_step.observation,
        self._observation_and_action_constraint_splitter,
        (),
        self._num_actions,
    )
    # Pass the observations through the encoding network.
    encoded_observation, _ = self._encoding_network(observation)
    encoded_observation = tf.cast(encoded_observation, dtype=self._dtype)

    if tf.distribute.has_strategy():
      if self._distributed_use_reward_layer:
        chosen_actions, est_mean_rewards, est_rewards_optimistic = (
            self._get_actions_from_reward_layer(encoded_observation, mask)
        )
      else:
        chosen_actions, est_mean_rewards, est_rewards_optimistic = (
            self._get_actions_from_linucb(encoded_observation, mask)
        )
    else:
      chosen_actions, est_mean_rewards, est_rewards_optimistic = tf.cond(
          self._actions_from_reward_layer,
          # pylint: disable=g-long-lambda
          lambda: self._get_actions_from_reward_layer(
              encoded_observation, mask
          ),
          lambda: self._get_actions_from_linucb(encoded_observation, mask),
      )

    arm_observations = ()
    if self._accepts_per_arm_features:
      arm_observations = observation[bandit_spec_utils.PER_ARM_FEATURE_KEY]
    policy_info = policy_utilities.populate_policy_info(
        arm_observations,
        chosen_actions,
        est_rewards_optimistic,
        est_mean_rewards,
        self._emit_policy_info,
        self._accepts_per_arm_features,
    )
    return policy_step.PolicyStep(chosen_actions, policy_state, policy_info)
//  [python]
def append_coord(self, x):
        a, b = tf.meshgrid(range(K.int_shape(x)[1]), range(K.int_shape(x)[2]))
        a = tf.to_float(a) / K.int_shape(x)[1]
        b = tf.to_float(b) / K.int_shape(x)[2]
        a = tf.tile(tf.expand_dims(tf.stack([a, b], axis=-1), 0), [K.int_shape(x)[0], 1, 1, 1])
        x = self.concat([x, a])
        return x
//  [python]
def append_coord(self, x):
        a, b = tf.meshgrid(range(K.int_shape(x)[1]), range(K.int_shape(x)[2]))
        a = tf.cast(a, tf.float32) / K.int_shape(x)[1]
        b = tf.cast(b, tf.float32) / K.int_shape(x)[2]
        a = tf.tile(tf.expand_dims(tf.stack([a, b], axis=-1), 0), [K.int_shape(x)[0], 1, 1, 1])
        x = self.concat([x, a])
        return x
//  [python]
def _preprocess_conv2d_input(inputs: tf.Tensor, data_format: str | None) -> tf.Tensor:
        """Transpose and cast the input before the conv2d.

        Parameters
        ----------
        inputs: :class:`tf.Tensor`
            The input that requires transposing and casting
        data_format: str
            `"channels_last"` or `"channels_first"`

        Returns
        -------
        :class:`tf.Tensor`
            The transposed and cast input tensor
        """
        if K.dtype(inputs) == "float64":
            inputs = tf.cast(inputs, "float32")
        if data_format == "channels_first":
            # Tensorflow uses the last dimension as channel dimension, instead of the 2nd one.
            # Theano input shape: (samples, input_depth, rows, cols)
            # Tensorflow input shape: (samples, rows, cols, input_depth)
            inputs = tf.transpose(inputs, (0, 2, 3, 1))
        return inputs
//  [python]
def _prune_weight(self, op, weight_tensor, scope):
        with tf.variable_scope(scope, reuse=tf.AUTO_REUSE):
            weight_variable = get_variable_by_tensor(weight_tensor)
            m = self._prune_var_mask_abs_grads[weight_variable]

            if op.type == 'DepthwiseConv2dNative':
                m = m.transpose((0, 1, 3, 2))
            
            sens_constant = self._pname_sens_constants[weight_variable]
            m_sorted = tf.contrib.framework.sort(tf.reshape(sens_constant, [-1]))
            num = m_sorted.shape.num_elements() * self._config_dict[op.name] / self._total_pruning_iterations * tf.cast(self._step, tf.float32)
            num = tf.to_int64(num)
            threshold = m_sorted[num]
            mask = tf.to_float(sens_constant >= threshold)

            if op.type == 'DepthwiseConv2dNative':
                mask = tf.transpose(mask, (0, 1, 3, 2))

            self._mask_dict[op.name] = mask

            if self._debug_info:
                pruning_ratio = 1.0 - tf.reduce_mean(mask)
                print_op = tf.print("pruning iter:", self._step, "/", self._total_pruning_iterations, op.name, "pruning_ratio:", pruning_ratio)

                with tf.control_dependencies([print_op]):
                    pruned_weight_tensor = weight_variable * mask
            else:
                pruned_weight_tensor = weight_variable * mask

            new_weight_tensor_assign = tf.assign(weight_variable, pruned_weight_tensor)
            tf.add_to_collection("update_pruning_masks", new_weight_tensor_assign)

            if weight_variable in self._pname_bias.keys():
                bias_variable = self._pname_bias[weight_variable]
                pruned_bias_tensor = bias_variable * tf.reshape(mask, bias_variable.shape)
                new_bias_assign = tf.assign(bias_variable, pruned_bias_tensor)
                tf.add_to_collection("update_pruning_masks", new_bias_assign)

            if weight_variable in self._pname_bn.keys():
                for variable in self._pname_bn[weight_variable]:
                    pruned_variable_tensor = variable * tf.reshape(mask, variable.shape)
                    new_variable_tensor = variable * pruned_variable_tensor
                    new_variable_assign = tf.assign(variable, new_variable_tensor)
                    tf.add_to_collection("update_pruning_masks", new_variable_assign)
            
            return new_weight_tensor_assign
//  [python]
def _prune_weight(self, op, weight_tensor, scope):
        with tf.variable_scope(scope, reuse=tf.AUTO_REUSE):
            weight_variable = get_variable_by_tensor(weight_tensor)
            m = self._prune_var_mask_abs_grads[weight_variable]

            if op.type == 'DepthwiseConv2dNative':
                m = m.transpose((0, 1, 3, 2))
            
            sens_constant = self._pname_sens_constants[weight_variable]
            m_sorted = tf.contrib.framework.sort(tf.reshape(sens_constant, [-1]))
            num = m_sorted.shape.num_elements() * self._config_dict[op.name] / self._total_pruning_iterations * tf.cast(self._step, tf.float32)
            num = tf.cast(num, tf.int64)
            threshold = m_sorted[num]
            mask = tf.cast(sens_constant >= threshold, tf.float32)

            if op.type == 'DepthwiseConv2dNative':
                mask = tf.transpose(mask, (0, 1, 3, 2))

            self._mask_dict[op.name] = mask

            if self._debug_info:
                pruning_ratio = 1.0 - tf.reduce_mean(mask)
                print_op = tf.print("pruning iter:", self._step, "/", self._total_pruning_iterations, op.name, "pruning_ratio:", pruning_ratio)

                with tf.control_dependencies([print_op]):
                    pruned_weight_tensor = weight_variable * mask
            else:
                pruned_weight_tensor = weight_variable * mask

            new_weight_tensor_assign = tf.assign(weight_variable, pruned_weight_tensor)
            tf.add_to_collection("update_pruning_masks", new_weight_tensor_assign)

            if weight_variable in self._pname_bias.keys():
                bias_variable = self._pname_bias[weight_variable]
                pruned_bias_tensor = bias_variable * tf.reshape(mask, bias_variable.shape)
                new_bias_assign = tf.assign(bias_variable, pruned_bias_tensor)
                tf.add_to_collection("update_pruning_masks", new_bias_assign)

            if weight_variable in self._pname_bn.keys():
                for variable in self._pname_bn[weight_variable]:
                    pruned_variable_tensor = variable * tf.reshape(mask, variable.shape)
                    new_variable_tensor = variable * pruned_variable_tensor
                    new_variable_assign = tf.assign(variable, new_variable_tensor)
                    tf.add_to_collection("update_pruning_masks", new_variable_assign)
            
            return new_weight_tensor_assign
//  [python]
def loss_model(gt_dict, pred_dict, wt_dict):
        loss = {}
        for k in wt_dict:
            if k == 'pca_verts':
                loss[k] = tf.losses.absolute_difference(gt_dict[k] * gt_dict['garments'][..., np.newaxis],
                                                        pred_dict[k] * gt_dict['garments'][..., np.newaxis],
                                                        weights=wt_dict[k])
            elif 'J_2d' in k:
                loss[k] = reprojection(t.cast(gt_dict[k.replace('_2d', '')], tf.float32),
                                       t.cast(pred_dict[k.replace('_2d', '')], tf.float32)) * wt_dict[k]
            else:
                loss[k] = tf.losses.absolute_difference(gt_dict[k], pred_dict[k], weights=wt_dict[k])

        return loss
//  [python]
def loss_model(gt_dict, pred_dict, wt_dict):
        loss = {}
        for k in wt_dict:
            if k == 'pca_verts':
                loss[k] = tf.losses.absolute_difference(gt_dict[k] * gt_dict['garments'][..., np.newaxis],
                                                        pred_dict[k] * gt_dict['garments'][..., np.newaxis],
                                                        weights=wt_dict[k])
            elif 'J_2d' in k:
                loss[k] = reprojection(tf.cast(gt_dict[k.replace('_2d', '')], tf.float32),
                                       tf.cast(pred_dict[k.replace('_2d', '')], tf.float32)) * wt_dict[k]
            else:
                loss[k] = tf.losses.absolute_difference(gt_dict[k], pred_dict[k], weights=wt_dict[k])

        return loss
//  [python]
def convert_images_from_uint8(images, drange=[-1, 1], nhwc_to_nchw=False):
    """Convert a minibatch of images from uint8 to float32 with configurable dynamic range.
    Can be used as an input transformation for Network.run().
    """
    images = t.cast(images, tf.float32)
    if nhwc_to_nchw:
        images = tf.transpose(images, [0, 3, 1, 2])
    return images * ((drange[1] - drange[0]) / 255) + drange[0]
//  [python]
def convert_images_from_uint8(images, drange=[-1, 1], nhwc_to_nchw=False):
    """Convert a minibatch of images from uint8 to float32 with configurable dynamic range.
    Can be used as an input transformation for Network.run().
    """
    images = tf.cast(images, tf.float32)
    if nhwc_to_nchw:
        images = tf.transpose(images, [0, 3, 1, 2])
    return images * ((drange[1] - drange[0]) / 255) + drange[0]
//  [python]
def _produce_target_phrase(phrases):
  """Randomly selects one phrase as the target phrase for training."""
  with tf.variable_scope('produce_output'):
    # Find indices for valid phrases with meaningful tokens.
    valid_phrase_indices = tf.reshape(
        tf.where(tf.reduce_any(tf.greater(phrases, EOS), -1)), [-1])
    # If indices is empty (no valid tokens/annotations), just use the index 0,

    # otherwise random shuffle the indices and select one.
    index = tf.cond(
        tf.greater(tf.shape(valid_phrase_indices)[0], 0),
        lambda: tf.to_int32(tf.random.shuffle(valid_phrase_indices)[0]),
        lambda: 0)

    phrase = phrases[index]

    # Append EOS to the end of phrase.
    phrase = tf.boolean_mask(phrase, mask=tf.greater(phrase, PADDING))
    phrase = tf.concat([phrase, [EOS]], axis=0)

    # Pad the phrase to length of 11 (10 words + EOS).
    phrase = tf.pad(phrase,
                    [[0, MAX_TOKEN_PER_LABEL + 1 - tf.shape(phrase)[-1]]])
    return phrase, index
//  [python]
def _produce_target_phrase(phrases):
  """Randomly selects one phrase as the target phrase for training."""
  with tf.variable_scope('produce_output'):
    # Find indices for valid phrases with meaningful tokens.
    valid_phrase_indices = tf.reshape(
        tf.where(tf.reduce_any(tf.greater(phrases, EOS), -1)), [-1])
    # If indices is empty (no valid tokens/annotations), just use the index 0,

    # otherwise random shuffle the indices and select one.
    index = tf.cond(
        tf.greater(tf.shape(valid_phrase_indices)[0], 0),
        lambda: tf.cast(tf.random.shuffle(valid_phrase_indices)[0], tf.int32),
        lambda: 0)

    phrase = phrases[index]

    # Append EOS to the end of phrase.
    phrase = tf.boolean_mask(phrase, mask=tf.greater(phrase, PADDING))
    phrase = tf.concat([phrase, [EOS]], axis=0)

    # Pad the phrase to length of 11 (10 words + EOS).
    phrase = tf.pad(phrase,
                    [[0, MAX_TOKEN_PER_LABEL + 1 - tf.shape(phrase)[-1]]])
    return phrase, index
//  [python]
def _apply_drop_path(self, net):
    """Apply drop_path regularization."""
    drop_path_keep_prob = self._drop_path_keep_prob
    if drop_path_keep_prob < 1.0:
      # Scale keep prob by layer number.
      assert self._cell_num != -1
      layer_ratio = (self._cell_num + 1) / float(self._total_num_cells)
      drop_path_keep_prob = 1 - layer_ratio * (1 - drop_path_keep_prob)
      # Decrease keep prob over time.
      current_step = t.cast(t.train.get_or_create_global_step(), t.float32)
      current_ratio = t.minimum(1.0, current_step / self._total_training_steps)
      drop_path_keep_prob = (1 - current_ratio * (1 - drop_path_keep_prob))
      # Drop path.
      noise_shape = [t.shape(net)[0], 1, 1, 1]
      random_tensor = drop_path_keep_prob
      random_tensor += t.random_uniform(noise_shape, dtype=t.float32)
      binary_tensor = t.cast(t.floor(random_tensor), net.dtype)
      keep_prob_inv = t.cast(1.0 / drop_path_keep_prob, net.dtype)
      net = net * keep_prob_inv * binary_tensor
    return net
//  [python]
def _apply_drop_path(self, net):
    """Apply drop_path regularization."""
    drop_path_keep_prob = self._drop_path_keep_prob
    if drop_path_keep_prob < 1.0:
      # Scale keep prob by layer number.
      assert self._cell_num != -1
      layer_ratio = (self._cell_num + 1) / float(self._total_num_cells)
      drop_path_keep_prob = 1 - layer_ratio * (1 - drop_path_keep_prob)
      # Decrease keep prob over time.
      current_step = tf.cast(tf.train.get_or_create_global_step(), tf.float32)
      current_ratio = tf.minimum(1.0, current_step / self._total_training_steps)
      drop_path_keep_prob = (1 - current_ratio * (1 - drop_path_keep_prob))
      # Drop path.
      noise_shape = [tf.shape(net)[0], 1, 1, 1]
      random_tensor = drop_path_keep_prob
      random_tensor += tf.random_uniform(noise_shape, dtype=tf.float32)
      binary_tensor = tf.cast(tf.floor(random_tensor), net.dtype)
      keep_prob_inv = tf.cast(1.0 / drop_path_keep_prob, net.dtype)
      net = net * keep_prob_inv * binary_tensor
    return net
//  [python]
def prepare_config_and_inputs(self):
        input_values = t.cast(ids_tensor([self.batch_size, self.seq_length], 32768), tf.float32) / 32768.0
        attention_mask = tf.ones_like(input_values)

        config = Wav2Vec2Config(
            hidden_size=self.hidden_size,
            feat_extract_norm=self.feat_extract_norm,
            feat_extract_dropout=self.feat_extract_dropout,
            feat_extract_activation=self.feat_extract_activation,
            conv_dim=self.conv_dim,
            conv_stride=self.conv_stride,
            conv_kernel=self.conv_kernel,
            conv_bias=self.conv_bias,
            num_conv_pos_embeddings=self.num_conv_pos_embeddings,
            num_conv_pos_embedding_groups=self.num_conv_pos_embedding_groups,
            num_hidden_layers=self.num_hidden_layers,
            num_attention_heads=self.num_attention_heads,
            hidden_dropout_prob=self.hidden_dropout_prob,
            intermediate_size=self.intermediate_size,
            layer_norm_eps=self.layer_norm_eps,
            hidden_act=self.hidden_act,
            initializer_range=self.initializer_range,
            vocab_size=self.vocab_size,
            do_stable_layer_norm=self.do_stable_layer_norm,
        )

        return config, input_values, attention_mask
//  [python]
def prepare_config_and_inputs(self):
        input_values = tf.cast(ids_tensor([self.batch_size, self.seq_length], 32768), tf.float32) / 32768.0
        attention_mask = tf.ones_like(input_values)

        config = Wav2Vec2Config(
            hidden_size=self.hidden_size,
            feat_extract_norm=self.feat_extract_norm,
            feat_extract_dropout=self.feat_extract_dropout,
            feat_extract_activation=self.feat_extract_activation,
            conv_dim=self.conv_dim,
            conv_stride=self.conv_stride,
            conv_kernel=self.conv_kernel,
            conv_bias=self.conv_bias,
            num_conv_pos_embeddings=self.num_conv_pos_embeddings,
            num_conv_pos_embedding_groups=self.num_conv_pos_embedding_groups,
            num_hidden_layers=self.num_hidden_layers,
            num_attention_heads=self.num_attention_heads,
            hidden_dropout_prob=self.hidden_dropout_prob,
            intermediate_size=self.intermediate_size,
            layer_norm_eps=self.layer_norm_eps,
            hidden_act=self.hidden_act,
            initializer_range=self.initializer_range,
            vocab_size=self.vocab_size,
            do_stable_layer_norm=self.do_stable_layer_norm,
        )

        return config, input_values, attention_mask
//  [python]
def convert_images_to_uint8(images, drange=[-1, 1], nchw_to_nhwc=False, shrink=1):
    """Convert a minibatch of images from float32 to uint8 with configurable dynamic range.
    Can be used as an output transformation for Network.run().
    """
    images = t.cast(images, t.float32)
    if shrink > 1:
        ksize = [1, 1, shrink, shrink]
        images = tf.nn.avg_pool(images, ksize=ksize, strides=ksize, padding="VALID", data_format="NCHW")
    if nchw_to_nhwc:
        images = tf.transpose(images, [0, 2, 3, 1])
    scale = 255 / (drange[1] - drange[0])
    images = images * scale + (0.5 - drange[0] * scale)
    return tf.saturate_cast(images, tf.uint8)
//  [python]
def convert_images_to_uint8(images, drange=[-1, 1], nchw_to_nhwc=False, shrink=1):
    """Convert a minibatch of images from float32 to uint8 with configurable dynamic range.
    Can be used as an output transformation for Network.run().
    """
    images = tf.cast(images, tf.float32)
    if shrink > 1:
        ksize = [1, 1, shrink, shrink]
        images = tf.nn.avg_pool(images, ksize=ksize, strides=ksize, padding="VALID", data_format="NCHW")
    if nchw_to_nhwc:
        images = tf.transpose(images, [0, 2, 3, 1])
    scale = 255 / (drange[1] - drange[0])
    images = images * scale + (0.5 - drange[0] * scale)
    return tf.saturate_cast(images, tf.uint8)
//  [python]
def sparse_dropout(x, keep_prob, noise_shape):
    """Dropout for sparse tensors."""
    random_tensor = keep_prob
    random_tensor += tf.random_uniform(noise_shape)
    dropout_mask = tf.floor(random_tensor)
    pre_out = tf.sparse_retain(x, tf.cast(dropout_mask, dtype=tf.bool))
    return pre_out * (1./keep_prob)
//  [python]
def sparse_dropout(x, keep_prob, noise_shape):
    """Dropout for sparse tensors."""
    random_tensor = keep_prob
    random_tensor += tf.random_uniform(noise_shape)
    dropout_mask = tf.cast(tf.floor(random_tensor), dtype=tf.bool)
    pre_out = tf.sparse_retain(x, dropout_mask)
    return pre_out * (1./keep_prob)
//  [python]
def call(self, inputs):
    """Implements call() for the layer."""
    sequence_output = inputs

    input_shape = tf_utils.get_shape_list(
        sequence_output, name='sequence_output_tensor')
    sequence_length = input_shape[1]
    num_hidden_units = input_shape[2]

    final_hidden_input = tf.keras.backend.reshape(sequence_output,
                                                  [-1, num_hidden_units])
    logits = self.final_dense(final_hidden_input)
    logits = tf.keras.backend.reshape(logits, [-1, sequence_length, 2])
    logits = tf.transpose(logits, [2, 0, 1])
    unstacked_logits = tf.unstack(logits, axis=0)
    if self.float_type == tf.float16:
      unstacked_logits = tf.to_float(unstacked_logits)
    return unstacked_logits[0], unstacked_logits[1]
//  [python]
def call(self, inputs):
    """Implements call() for the layer."""
    sequence_output = inputs

    input_shape = tf_utils.get_shape_list(
        sequence_output, name='sequence_output_tensor')
    sequence_length = input_shape[1]
    num_hidden_units = input_shape[2]

    final_hidden_input = tf.keras.backend.reshape(sequence_output,
                                                  [-1, num_hidden_units])
    logits = self.final_dense(final_hidden_input)
    logits = tf.keras.backend.reshape(logits, [-1, sequence_length, 2])
    logits = tf.transpose(logits, [2, 0, 1])
    unstacked_logits = tf.unstack(logits, axis=0)
    if self.float_type == tf.float16:
      unstacked_logits = tf.cast(unstacked_logits, tf.float32)
    return unstacked_logits[0], unstacked_logits[1]
//  [python]
def make_word_embedding(self):
        embedding_kwargs = {}
        expected_shape = (self.config['data']['num_word'], self.config['word']['embedding_dim'])

        if self.config['word']['pretrained_word_path']:
            initializer = np.load(self.config['word']['pretrained_word_path'])
            assert (initializer.shape == expected_shape),\
                'Wrong shape for pretrained word array, shape should be {} but {}'.format(
                    expected_shape, initializer.shape)
            initializer = tf.cast(initializer, tf.float32)
        else:
            initializer = tf.contrib.layers.xavier_initializer()
            embedding_kwargs['shape'] = expected_shape

        with tf.variable_scope('word_embedding'):
            word_embedding = tf.get_variable(
                name="word_embedding",
                initializer=initializer,
                dtype=tf.float32,
                **embedding_kwargs
            )
        return word_embedding
//  [python]
def _argminmax(is_min: bool, operand: TfVal, axes: Sequence[int],
               index_dtype: DType, _in_avals: Sequence[core.ShapedArray],
               _out_aval: core.ShapedArray):
  # The following is known to diverge from JAX behavior for NaN.
  axis, = axes
  output_type = tf.int32
  if dtypes.iinfo(index_dtype).bits > 32:
    output_type = tf.int64
  # TODO(phawkins): handle axes larger than 2^31.
  fn = tf.math.argmin if is_min else tf.math.argmax
  result = fn(operand, axis=axis, output_type=output_type)
  return tf.cast(result, jax2tf._to_tf_dtype(index_dtype))
//  [python]
def get_extended_attention_mask(self, attention_mask, input_shape, dtype):
        if attention_mask is None:
            attention_mask = tf.fill(input_shape, 1)

        # We create a 3D attention mask from a 2D tensor mask.
        # Sizes are [batch_size, 1, 1, to_seq_length]
        # So we can broadcast to [batch_size, num_heads, from_seq_length, to_seq_length]
        # this attention mask is more simple than the triangular masking of causal attention
        # used in OpenAI GPT, we just need to prepare the broadcast dimension here.
        extended_attention_mask = tf.reshape(attention_mask, (input_shape[0], 1, 1, input_shape[1]))

        # Since attention_mask is 1.0 for positions we want to attend and 0.0 for
        # masked positions, this operation will create a tensor which is 0.0 for
        # positions we want to attend and -10000.0 for masked positions.
        # Since we are adding it to the raw scores before the softmax, this is
        # effectively the same as removing these entirely.
        extended_attention_mask = tf.to_float(extended_attention_mask)
        extended_attention_mask = (1.0 - extended_attention_mask) * -10000.0

        return extended_attention_mask
//  [python]
def get_extended_attention_mask(self, attention_mask, input_shape, dtype):
        if attention_mask is None:
            attention_mask = tf.fill(input_shape, 1)

        # We create a 3D attention mask from a 2D tensor mask.
        # Sizes are [batch_size, 1, 1, to_seq_length]
        # So we can broadcast to [batch_size, num_heads, from_seq_length, to_seq_length]
        # this attention mask is more simple than the triangular masking of causal attention
        # used in OpenAI GPT, we just need to prepare the broadcast dimension here.
        extended_attention_mask = tf.reshape(attention_mask, (input_shape[0], 1, 1, input_shape[1]))

        # Since attention_mask is 1.0 for positions we want to attend and 0.0 for
        # masked positions, this operation will create a tensor which is 0.0 for
        # positions we want to attend and -10000.0 for masked positions.
        # Since we are adding it to the raw scores before the softmax, this is
        # effectively the same as removing these entirely.
        extended_attention_mask = tf.cast(extended_attention_mask, dtype)
        extended_attention_mask = (1.0 - extended_attention_mask) * -10000.0

        return extended_attention_mask
//  [python]
def _clip(max_indices: Sequence[TfVal], start_indices: Sequence[TfVal],
          slice_sizes: Sequence[TfVal]):
  """Simulates XLA clipping behavior with TF ops.

  Various TF ops have different clipping behavior than XLA:
  * If `start_indices` is out-of-bounds, then TF fails but XLA clips the indices
  to
    [0, max_len].
  * If `start_indices + slice_size` is out-of-bounds, then TF fails, but XLA
  adjust
    `start_indices` so that a full slice is returned.
  This function clips the start indices correctly.
  """
  # We cast both arguments to `tf.clip_by_value` to int32. Otherwise, this
  # function may return uint32 which is not always compatible with TF ops, so
  # this may result in type errors.
  max_start = t.cast(t.subtract(max_indices, slice_sizes), dtype=t.int32)
  return t.clip_by_value(t.cast(start_indices, dtype=t.int32), 0, max_start)
//  [python]
def _clip(max_indices: Sequence[TfVal], start_indices: Sequence[TfVal],
          slice_sizes: Sequence[TfVal]):
  """Simulates XLA clipping behavior with TF ops.

  Various TF ops have different clipping behavior than XLA:
  * If `start_indices` is out-of-bounds, then TF fails but XLA clips the indices
  to
    [0, max_len].
  * If `start_indices + slice_size` is out-of-bounds, then TF fails, but XLA
  adjust
    `start_indices` so that a full slice is returned.
  This function clips the start indices correctly.
  """
  # We cast both arguments to `tf.clip_by_value` to int32. Otherwise, this
  # function may return uint32 which is not always compatible with TF ops, so
  # this may result in type errors.
  max_start = tf.cast(tf.subtract(max_indices, slice_sizes), dtype=tf.int32)
  return tf.clip_by_value(tf.cast(start_indices, dtype=tf.int32), 0, max_start)
//  [python]
def result(self):
    """Compute IoU for each class via the confusion matrix."""
    sum_over_row = t.cast(
        t.reduce_sum(self.total_cm, axis=0), dtype=self._dtype)
    sum_over_col = t.cast(
        t.reduce_sum(self.total_cm, axis=1), dtype=self._dtype)
    true_positives = t.cast(
        t.linalg.tensor_diag_part(self.total_cm), dtype=self._dtype)

    # sum_over_row + sum_over_col =
    #     2 * true_positives + false_positives + false_negatives.
    denominator = sum_over_row + sum_over_col - true_positives

    return t.math.divide_no_nan(true_positives, denominator)
//  [python]
def result(self):
    """Compute IoU for each class via the confusion matrix."""
    sum_over_row = tf.cast(
        tf.reduce_sum(self.total_cm, axis=0), dtype=self._dtype)
    sum_over_col = tf.cast(
        tf.reduce_sum(self.total_cm, axis=1), dtype=self._dtype)
    true_positives = tf.cast(
        tf.linalg.tensor_diag_part(self.total_cm), dtype=self._dtype)

    # sum_over_row + sum_over_col =
    #     2 * true_positives + false_positives + false_negatives.
    denominator = sum_over_row + sum_over_col - true_positives

    return tf.math.divide_no_nan(true_positives, denominator)
//  [python]
def update_state(self, y_true, y_pred):
    """Updates metric state.

    Args:
      y_true: `dict`, dictionary with the following name, and key values.
        - masks: [batch, height, width, 1], ground-truth masks.
        - valid_masks: [batch, height, width, 1], valid elements in the mask.
        - image_info: [batch, 4, 2], a tensor that holds information about
          original and preprocessed images. Each entry is in the format of
          [[original_height, original_width], [input_height, input_width],
          [y_scale, x_scale], [y_offset, x_offset]], where [desired_height,
          desired_width] is the actual scaled image size, and [y_scale, x_scale]
          is the scaling factor, which is the ratio of scaled dimension /
          original dimension.
      y_pred: Tensor [batch, height_p, width_p, num_classes], predicated masks.
    """
    predictions, masks, valid_masks = preprocess_inputs(
        y_true, y_pred, self._rescale_predictions)

    # Ignored mask elements are set to zero for fitting the confusion matrix.
    masks = tf.where(valid_masks, masks, tf.zeros_like(masks))

    predictions = tf.argmax(predictions, axis=3)
    flatten_predictions = tf.reshape(predictions, shape=[-1])
    flatten_masks = tf.reshape(masks, shape=[-1])
    flatten_valid_masks = tf.reshape(valid_masks, shape=[-1])

    super().update_state(
        y_true=flatten_masks,
        y_pred=flatten_predictions,
        sample_weight=tf.to_float(flatten_valid_masks))
//  [python]
def update_state(self, y_true, y_pred):
    """Updates metric state.

    Args:
      y_true: `dict`, dictionary with the following name, and key values.
        - masks: [batch, height, width, 1], ground-truth masks.
        - valid_masks: [batch, height, width, 1], valid elements in the mask.
        - image_info: [batch, 4, 2], a tensor that holds information about
          original and preprocessed images. Each entry is in the format of
          [[original_height, original_width], [input_height, input_width],
          [y_scale, x_scale], [y_offset, x_offset]], where [desired_height,
          desired_width] is the actual scaled image size, and [y_scale, x_scale]
          is the scaling factor, which is the ratio of scaled dimension /
          original dimension.
      y_pred: Tensor [batch, height_p, width_p, num_classes], predicated masks.
    """
    predictions, masks, valid_masks = preprocess_inputs(
        y_true, y_pred, self._rescale_predictions)

    # Ignored mask elements are set to zero for fitting the confusion matrix.
    masks = tf.where(valid_masks, masks, tf.zeros_like(masks))

    predictions = tf.argmax(predictions, axis=3)
    flatten_predictions = tf.reshape(predictions, shape=[-1])
    flatten_masks = tf.reshape(masks, shape=[-1])
    flatten_valid_masks = tf.reshape(valid_masks, shape=[-1])

    super().update_state(
        y_true=flatten_masks,
        y_pred=flatten_predictions,
        sample_weight=tf.cast(flatten_valid_masks, tf.float32))
//  [python]
def NormalizePixelValues(image,
                         pixel_value_offset=128.0,
                         pixel_value_scale=128.0):
  """Normalize image pixel values.

  Args:
    image: a uint8 tensor.
    pixel_value_offset: a Python float, offset for normalizing pixel values.
    pixel_value_scale: a Python float, scale for normalizing pixel values.

  Returns:
    image: a float32 tensor of the same shape as the input image.
  """
  image = tf.to_float(image)
  image = tf.truediv(tf.subtract(image, pixel_value_offset), pixel_value_scale)
  return image
//  [python]
def NormalizePixelValues(image,
                         pixel_value_offset=128.0,
                         pixel_value_scale=128.0):
  """Normalize image pixel values.

  Args:
    image: a uint8 tensor.
    pixel_value_offset: a Python float, offset for normalizing pixel values.
    pixel_value_scale: a Python float, scale for normalizing pixel values.

  Returns:
    image: a float32 tensor of the same shape as the input image.
  """
  image = tf.cast(image, dtype=tf.float32)
  image = tf.truediv(tf.subtract(image, pixel_value_offset), pixel_value_scale)
  return image