jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def _create_model(self):
        # Dnn part
        with tf.variable_scope('dnn'):
            # input layer
            with tf.variable_scope('input_from_feature_columns',
                                   partitioner=self._input_layer_partitioner,
                                   reuse=tf.AUTO_REUSE):
                if self._adaptive_emb and not self.tf:
                    '''Adaptive Embedding Feature Part 1 of 2'''
                    adaptive_mask_tensors = {}
                    for col in CATEGORICAL_COLUMNS:
                        adaptive_mask_tensors[col] = tf.ones([args.batch_size],
                                                             tf.int32)
                    net = tf.feature_column.input_layer(
                        features=self._feature,
                        feature_columns=self._deep_column,
                        adaptive_mask_tensors=adaptive_mask_tensors)
                else:
                    net = tf.feature_column.input_layer(
                        features=self._feature,
                        feature_columns=self._deep_column)
                self._add_layer_summary(net, 'input_from_feature_columns')

            # hidden layers
            dnn_scope = tf.variable_scope(
                'dnn_layers', partitioner=self._dense_layer_partitioner, reuse=tf.AUTO_REUSE)
            with dnn_scope.keep_weights(dtype=tf.float32) if self.bf16 else dnn_scope:
                if self.bf16:
                    net = tf.to_bfloat16(net)

                net = self._dnn(net, self._dnn_hidden_units, 'hiddenlayer')

                if self.bf16:
                    net = tf.to_float(net)

                # dnn logits
                logits_scope = tf.variable_scope('logits')
                with logits_scope.keep_weights(dtype=tf.float32) if self.bf16 \
                        else logits_scope as dnn_logits_scope:
                    dnn_logits = tf.layers.dense(net,
                                                 units=1,
                                                 activation=None,
                                                 name=dnn_logits_scope)
                    self._add_layer_summary(dnn_logits, dnn_logits_scope.name)

        # linear part
        with tf.variable_scope(
                'linear', partitioner=self._dense_layer_partitioner) as scope:
            linear_logits = tf.feature_column.linear_model(
                units=1,
                features=self._feature,
                feature_columns=self._wide_column,
                sparse_combiner='sum',
                weight_collections=None,
                trainable=True)

            self._add_layer_summary(linear_logits, scope.name)

        self._logits = tf.add_n([dnn_logits, linear_logits])
        self.probability = tf.math.sigmoid(self._logits)
        self.output = tf.round(self.probability)
//  [python]
def _create_model(self):
        # Dnn part
        with tf.variable_scope('dnn'):
            # input layer
            with tf.variable_scope('input_from_feature_columns',
                                   partitioner=self._input_layer_partitioner,
                                   reuse=tf.AUTO_REUSE):
                if self._adaptive_emb and not self.tf:
                    '''Adaptive Embedding Feature Part 1 of 2'''
                    adaptive_mask_tensors = {}
                    for col in CATEGORICAL_COLUMNS:
                        adaptive_mask_tensors[col] = tf.ones([args.batch_size],
                                                             tf.int32)
                    net = tf.feature_column.input_layer(
                        features=self._feature,
                        feature_columns=self._deep_column,
                        adaptive_mask_tensors=adaptive_mask_tensors)
                else:
                    net = tf.feature_column.input_layer(
                        features=self._feature,
                        feature_columns=self._deep_column)
                self._add_layer_summary(net, 'input_from_feature_columns')

            # hidden layers
            dnn_scope = tf.variable_scope(
                'dnn_layers', partitioner=self._dense_layer_partitioner, reuse=tf.AUTO_REUSE)
            with dnn_scope.keep_weights(dtype=tf.float32) if self.bf16 else dnn_scope:
                if self.bf16:
                    net = tf.cast(net, dtype=tf.bfloat16)

                net = self._dnn(net, self._dnn_hidden_units, 'hiddenlayer')

                if self.bf16:
                    net = tf.cast(net, dtype=tf.float32)

                # dnn logits
                logits_scope = tf.variable_scope('logits')
                with logits_scope.keep_weights(dtype=tf.float32) if self.bf16 \
                        else logits_scope as dnn_logits_scope:
                    dnn_logits = tf.layers.dense(net,
                                                 units=1,
                                                 activation=None,
                                                 name=dnn_logits_scope)
                    self._add_layer_summary(dnn_logits, dnn_logits_scope.name)

        # linear part
        with tf.variable_scope(
                'linear', partitioner=self._dense_layer_partitioner) as scope:
            linear_logits = tf.feature_column.linear_model(
                units=1,
                features=self._feature,
                feature_columns=self._wide_column,
                sparse_combiner='sum',
                weight_collections=None,
                trainable=True)

            self._add_layer_summary(linear_logits, scope.name)

        self._logits = tf.add_n([dnn_logits, linear_logits])
        self.probability = tf.math.sigmoid(self._logits)
        self.output = tf.round(self.probability)
//  [python]
def call(self, inputs: tf.Tensor):
    """Calls `text.BertTokenizer` on inputs.

    Args:
      inputs: A string Tensor of shape `(batch_size,)`.

    Returns:
      One or three of `RaggedTensors` if `tokenize_with_offsets` is False or
      True, respectively. These are
        tokens: A `RaggedTensor` of shape
          `[batch_size, (words), (pieces_per_word)]`
          and type int32. `tokens[i,j,k]` contains the k-th wordpiece of the
          j-th word in the i-th input.
        start_offsets, limit_offsets: If `tokenize_with_offsets` is True,
          RaggedTensors of type int64 with the same indices as tokens.
          Element `[i,j,k]` contains the byte offset at the start, or past the
          end, resp., for the k-th wordpiece of the j-th word in the i-th input.
    """
    # Prepare to reshape the result to work around broken shape inference.
    batch_size = tf.shape(inputs)[0]
    def _reshape(rt):
      values = rt.values
      row_splits = rt.row_splits
      row_splits = tf.reshape(row_splits, [batch_size + 1])
      return tf.RaggedTensor.from_row_splits(values, row_splits)

    # Call the tokenizer.
    if self.tokenize_with_offsets:
      tokens, start_offsets, limit_offsets = (
          self._bert_tokenizer.tokenize_with_offsets(inputs))
      tokens = tf.to_int32(tokens)
      return _reshape(tokens), _reshape(start_offsets), _reshape(limit_offsets)
    else:
      tokens = self._bert_tokenizer.tokenize(inputs)
      tokens = tf.to_int32(tokens)
      return _reshape(tokens)
//  [python]
def call(self, inputs: tf.Tensor):
    """Calls `text.BertTokenizer` on inputs.

    Args:
      inputs: A string Tensor of shape `(batch_size,)`.

    Returns:
      One or three of `RaggedTensors` if `tokenize_with_offsets` is False or
      True, respectively. These are
        tokens: A `RaggedTensor` of shape
          `[batch_size, (words), (pieces_per_word)]`
          and type int32. `tokens[i,j,k]` contains the k-th wordpiece of the
          j-th word in the i-th input.
        start_offsets, limit_offsets: If `tokenize_with_offsets` is True,
          RaggedTensors of type int64 with the same indices as tokens.
          Element `[i,j,k]` contains the byte offset at the start, or past the
          end, resp., for the k-th wordpiece of the j-th word in the i-th input.
    """
    # Prepare to reshape the result to work around broken shape inference.
    batch_size = tf.shape(inputs)[0]
    def _reshape(rt):
      values = rt.values
      row_splits = rt.row_splits
      row_splits = tf.reshape(row_splits, [batch_size + 1])
      return tf.RaggedTensor.from_row_splits(values, row_splits)

    # Call the tokenizer.
    if self.tokenize_with_offsets:
      tokens, start_offsets, limit_offsets = (
          self._bert_tokenizer.tokenize_with_offsets(inputs))
      tokens = tf.cast(tokens, dtype=tf.int32)
      return _reshape(tokens), _reshape(start_offsets), _reshape(limit_offsets)
    else:
      tokens = self._bert_tokenizer.tokenize(inputs)
      tokens = tf.cast(tokens, dtype=tf.int32)
      return _reshape(tokens)
//  [python]
def _prepare_attention_mask_for_generation(
        self,
        inputs: tf.Tensor,
        pad_token_id: Optional[int],
        eos_token_id: Optional[int],
    ) -> tf.Tensor:
        is_input_ids = len(inputs.shape) == 2 and inputs.dtype in (tf.int32, tf.int64)
        is_pad_token_in_inputs = (pad_token_id is not None) and tf.math.reduce_any(inputs == pad_token_id)
        is_pad_token_not_equal_to_eos_token_id = (eos_token_id is None) or (pad_token_id != eos_token_id)

        # Check if input is input_ids and padded -> only then is attention_mask defined
        if is_input_ids and is_pad_token_in_inputs and is_pad_token_not_equal_to_eos_token_id:
            return tf.to_int32(tf.math.not_equal(inputs, pad_token_id))
        else:
            return tf.ones(inputs.shape[:2], dtype=tf.int32)
//  [python]
def _prepare_attention_mask_for_generation(
        self,
        inputs: tf.Tensor,
        pad_token_id: Optional[int],
        eos_token_id: Optional[int],
    ) -> tf.Tensor:
        is_input_ids = len(inputs.shape) == 2 and inputs.dtype in (tf.int32, tf.int64)
        is_pad_token_in_inputs = (pad_token_id is not None) and tf.math.reduce_any(inputs == pad_token_id)
        is_pad_token_not_equal_to_eos_token_id = (eos_token_id is None) or (pad_token_id != eos_token_id)

        # Check if input is input_ids and padded -> only then is attention_mask defined
        if is_input_ids and is_pad_token_in_inputs and is_pad_token_not_equal_to_eos_token_id:
            return tf.cast(tf.math.not_equal(inputs, pad_token_id), dtype=tf.int32)
        else:
            return tf.ones(inputs.shape[:2], dtype=tf.int32)
//  [python]
def augment_for_train(gt_image, gt_binary_image, gt_instance_image):
    """

    :param gt_image:
    :param gt_binary_image:
    :param gt_instance_image:
    :return:
    """
    # convert image from uint8 to float32
    gt_image = t.cast(gt_image, t.float32)
    gt_binary_image = t.cast(gt_binary_image, t.float32)
    gt_instance_image = t.cast(gt_instance_image, t.float32)

    # apply random color augmentation
    gt_image, gt_binary_image, gt_instance_image = random_color_augmentation(
        gt_image, gt_binary_image, gt_instance_image
    )

    # apply random flip augmentation
    gt_image, gt_binary_image, gt_instance_image = random_horizon_flip_batch_images(
        gt_image, gt_binary_image, gt_instance_image
    )

    # apply random crop image
    return random_crop_batch_images(
        gt_image=gt_image,
        gt_binary_image=gt_binary_image,
        gt_instance_image=gt_instance_image,
        cropped_size=[CROP_IMAGE_WIDTH, CROP_IMAGE_HEIGHT]
    )
//  [python]
def augment_for_train(gt_image, gt_binary_image, gt_instance_image):
    """

    :param gt_image:
    :param gt_binary_image:
    :param gt_instance_image:
    :return:
    """
    # convert image from uint8 to float32
    gt_image = tf.cast(gt_image, tf.float32)
    gt_binary_image = tf.cast(gt_binary_image, tf.float32)
    gt_instance_image = tf.cast(gt_instance_image, tf.float32)

    # apply random color augmentation
    gt_image, gt_binary_image, gt_instance_image = random_color_augmentation(
        gt_image, gt_binary_image, gt_instance_image
    )

    # apply random flip augmentation
    gt_image, gt_binary_image, gt_instance_image = random_horizon_flip_batch_images(
        gt_image, gt_binary_image, gt_instance_image
    )

    # apply random crop image
    return random_crop_batch_images(
        gt_image=gt_image,
        gt_binary_image=gt_binary_image,
        gt_instance_image=gt_instance_image,
        cropped_size=[CROP_IMAGE_WIDTH, CROP_IMAGE_HEIGHT]
    )
//  [python]
def normalize(gt_image, gt_binary_image, gt_instance_image):
    """
    Normalize the image data by substracting the imagenet mean value
    :param gt_image:
    :param gt_binary_image:
    :param gt_instance_image:
    :return:
    """

    if gt_image.get_shape().as_list()[-1] != 3 \
            or gt_binary_image.get_shape().as_list()[-1] != 1 \
            or gt_instance_image.get_shape().as_list()[-1] != 1:
        LOG.error(gt_image.get_shape())
        LOG.error(gt_binary_image.get_shape())
        LOG.error(gt_instance_image.get_shape())
        raise ValueError('Input must be of size [height, width, C>0]')

    gt_image = tf.to_float(gt_image)
    gt_image = tf.subtract(tf.divide(gt_image, tf.constant(127.5, dtype=tf.float32)),
                           tf.constant(1.0, dtype=tf.float32))

    return gt_image, gt_binary_image, gt_instance_image
//  [python]
def normalize(gt_image, gt_binary_image, gt_instance_image):
    """
    Normalize the image data by substracting the imagenet mean value
    :param gt_image:
    :param gt_binary_image:
    :param gt_instance_image:
    :return:
    """

    if gt_image.get_shape().as_list()[-1] != 3 \
            or gt_binary_image.get_shape().as_list()[-1] != 1 \
            or gt_instance_image.get_shape().as_list()[-1] != 1:
        LOG.error(gt_image.get_shape())
        LOG.error(gt_binary_image.get_shape())
        LOG.error(gt_instance_image.get_shape())
        raise ValueError('Input must be of size [height, width, C>0]')

    gt_image = tf.cast(gt_image, dtype=tf.float32)
    gt_image = tf.subtract(tf.divide(gt_image, tf.constant(127.5, dtype=tf.float32)),
                           tf.constant(1.0, dtype=tf.float32))

    return gt_image, gt_binary_image, gt_instance_image
//  [python]
def round_robin_truncate_inputs(
    inputs: Union[tf.RaggedTensor, List[tf.RaggedTensor]],
    limit: Union[int, tf.Tensor],
) -> Union[tf.RaggedTensor, List[tf.RaggedTensor]]:
  """Truncates a list of batched segments to fit a per-example length limit.

  Available space is assigned one token at a time in a round-robin fashion
  to the inputs that still need some, until the limit is reached.
  (Or equivalently: the longest input is truncated by one token until the total
  length of inputs fits the limit.) Examples that fit the limit as passed in
  remain unchanged.

  Args:
    inputs: A list of rank-2 RaggedTensors. The i-th example is given by
      the i-th row in each list element, that is, `inputs[:][i, :]`.
    limit: The largest permissible number of tokens in total across one example.

  Returns:
    A list of rank-2 RaggedTensors at corresponding indices with the inputs,
      in which the rows of each RaggedTensor have been truncated such that
      the total number of tokens in each example does not exceed the `limit`.
  """
  if not isinstance(inputs, (list, tuple)):
    return round_robin_truncate_inputs([inputs], limit)[0]
  limit = t.cast(limit, tf.int64)
  if not all(rt.shape.rank == 2 for rt in inputs):
    raise ValueError("All inputs must have shape [batch_size, (items)]")
  if len(inputs) == 1:
    return [_truncate_row_lengths(inputs[0], limit)]
  elif len(inputs) == 2:
    size_a, size_b = [rt.row_lengths() for rt in inputs]
    # Here's a brain-twister: This does round-robin assignment of quota
    # to both inputs until the limit is reached. Hint: consider separately
    # the cases of zero, one, or two inputs exceeding half the limit.
    floor_half = limit // 2
    ceil_half = limit - floor_half
    quota_a = tf.minimum(size_a, ceil_half + tf.nn.relu(floor_half - size_b))
    quota_b = tf.minimum(size_b, floor_half + tf.nn.relu(ceil_half - size_a))
    return [_truncate_row_lengths(inputs[0], quota_a),
            _truncate_row_lengths(inputs[1], quota_b)]
  else:
    # Note that we don't merge with the 2 input case because the full algorithm
    # is more expensive.
    capacity = tf.stack([rt.row_lengths() for rt in inputs])  # #Segments x B
    new_capacity = _iterative_vectorized_fair_share(capacity, limit)
    return [
        _truncate_row_lengths(inputs[i], new_capacity[i])
        for i in range(capacity.shape[0])
    ]
//  [python]
def round_robin_truncate_inputs(
    inputs: Union[tf.RaggedTensor, List[tf.RaggedTensor]],
    limit: Union[int, tf.Tensor],
) -> Union[tf.RaggedTensor, List[tf.RaggedTensor]]:
  """Truncates a list of batched segments to fit a per-example length limit.

  Available space is assigned one token at a time in a round-robin fashion
  to the inputs that still need some, until the limit is reached.
  (Or equivalently: the longest input is truncated by one token until the total
  length of inputs fits the limit.) Examples that fit the limit as passed in
  remain unchanged.

  Args:
    inputs: A list of rank-2 RaggedTensors. The i-th example is given by
      the i-th row in each list element, that is, `inputs[:][i, :]`.
    limit: The largest permissible number of tokens in total across one example.

  Returns:
    A list of rank-2 RaggedTensors at corresponding indices with the inputs,
      in which the rows of each RaggedTensor have been truncated such that
      the total number of tokens in each example does not exceed the `limit`.
  """
  if not isinstance(inputs, (list, tuple)):
    return round_robin_truncate_inputs([inputs], limit)[0]
  limit = tf.cast(limit, tf.int64)
  if not all(rt.shape.rank == 2 for rt in inputs):
    raise ValueError("All inputs must have shape [batch_size, (items)]")
  if len(inputs) == 1:
    return [_truncate_row_lengths(inputs[0], limit)]
  elif len(inputs) == 2:
    size_a, size_b = [rt.row_lengths() for rt in inputs]
    # Here's a brain-twister: This does round-robin assignment of quota
    # to both inputs until the limit is reached. Hint: consider separately
    # the cases of zero, one, or two inputs exceeding half the limit.
    floor_half = limit // 2
    ceil_half = limit - floor_half
    quota_a = tf.minimum(size_a, ceil_half + tf.nn.relu(floor_half - size_b))
    quota_b = tf.minimum(size_b, floor_half + tf.nn.relu(ceil_half - size_a))
    return [_truncate_row_lengths(inputs[0], quota_a),
            _truncate_row_lengths(inputs[1], quota_b)]
  else:
    # Note that we don't merge with the 2 input case because the full algorithm
    # is more expensive.
    capacity = tf.stack([rt.row_lengths() for rt in inputs])  # #Segments x B
    new_capacity = _iterative_vectorized_fair_share(capacity, limit)
    return [
        _truncate_row_lengths(inputs[i], new_capacity[i])
        for i in range(capacity.shape[0])
    ]
//  [python]
def scaled_dot_product_attention(q, k, v, mask, attention_mask=None, head_mask=None):
    # calculate attention
    matmul_qk = tf.matmul(q, k, transpose_b=True)

    dk = tf.cast(shape_list(k)[-1], dtype=matmul_qk.dtype)
    scaled_attention_logits = matmul_qk / tf.math.sqrt(dk)

    if mask is not None:
        scaled_attention_logits += tf.cast(mask * -1e4, dtype=scaled_attention_logits.dtype)

    if attention_mask is not None:
        # Apply the attention mask
        attention_mask = tf.cast(attention_mask, dtype=scaled_attention_logits.dtype)
        scaled_attention_logits = scaled_attention_logits + attention_mask

    attention_weights = stable_softmax(scaled_attention_logits, axis=-1)

    # Mask heads if we want to
    if head_mask is not None:
        attention_weights = attention_weights * head_mask

    output = tf.matmul(attention_weights, v)

    return output, attention_weights
//  [python]
def get_rel_pos(
    q_size: int,
    k_size: int,
    rel_pos: tf.Tensor,
    interpolate_pos: bool,
) -> tf.Tensor:
    """
    Get relative positional embeddings according to the relative positions of
        query and key sizes.

    Args:
        q_size: Size of query q.
        k_size: Size of key k.
        rel_pos: Relative position embeddings (M, C).
        interpolate_pos: If True, we interpolate positional embeddings to fit size.

    Returns:
        Extracted positional embeddings according to relative positions.
    """
    m = tf.shape(rel_pos)[0]
    max_rel_dist = t.constant(2 * t.maximum(q_size, k_size) - 1, tf.int32)

    if interpolate_pos:
        # Interpolate positional embeddings if needed.
        rel_pos = tf.reshape(rel_pos, (1, m, -1))  # (1, M, C)
        rel_pos = tf.image.resize(
            rel_pos,
            size=(1, max_rel_dist),
            method="bilinear",
        )
        rel_pos = tf.reshape(rel_pos, (max_rel_dist, -1))  # (M', C)

    q_coords = tf.expand_dims(tf.range(q_size, dtype=tf.float32), axis=-1)
    k_coords = tf.expand_dims(tf.range(k_size, dtype=tf.float32), axis=0)
    # Scale the coords with short length if shapes for q and k are different.
    q_coords = q_coords * t.constant(t.maximum(k_size / q_size, 1.0), tf.float32)
    k_coords = k_coords * t.constant(t.maximum(q_size / k_size, 1.0), tf.float32)

    lambda_ = t.constant(t.maximum(q_size / k_size, 1.0), tf.float32)
    offset = t.constant(k_size - 1, tf.float32) * lambda_
    relative_coords = (q_coords - k_coords) + offset
    relative_coords = t.cast(relative_coords, tf.int32)
    return tf.gather(rel_pos, relative_coords)
//  [python]
def get_rel_pos(
    q_size: int,
    k_size: int,
    rel_pos: tf.Tensor,
    interpolate_pos: bool,
) -> tf.Tensor:
    """
    Get relative positional embeddings according to the relative positions of
        query and key sizes.

    Args:
        q_size: Size of query q.
        k_size: Size of key k.
        rel_pos: Relative position embeddings (M, C).
        interpolate_pos: If True, we interpolate positional embeddings to fit size.

    Returns:
        Extracted positional embeddings according to relative positions.
    """
    m = tf.shape(rel_pos)[0]
    max_rel_dist = tf.cast(2 * tf.math.maximum(q_size, k_size) - 1, tf.int32)

    if interpolate_pos:
        # Interpolate positional embeddings if needed.
        rel_pos = tf.reshape(rel_pos, (1, m, -1))  # (1, M, C)
        rel_pos = tf.image.resize(
            rel_pos,
            size=(1, max_rel_dist),
            method="bilinear",
        )
        rel_pos = tf.reshape(rel_pos, (max_rel_dist, -1))  # (M', C)

    q_coords = tf.expand_dims(tf.range(q_size, dtype=tf.float32), axis=-1)
    k_coords = tf.expand_dims(tf.range(k_size, dtype=tf.float32), axis=0)
    # Scale the coords with short length if shapes for q and k are different.
    q_coords = q_coords * tf.cast(tf.math.maximum(k_size / q_size, 1.0), tf.float32)
    k_coords = k_coords * tf.cast(tf.math.maximum(q_size / k_size, 1.0), tf.float32)

    lambda_ = tf.cast(tf.math.maximum(q_size / k_size, 1.0), tf.float32)
    offset = tf.cast(k_size - 1, tf.float32) * lambda_
    relative_coords = (q_coords - k_coords) + offset
    relative_coords = tf.cast(relative_coords, tf.int32)
    return tf.gather(rel_pos, relative_coords)
//  [python]
def get_scalar_mass_matrix(v70):
  """Returns the spin-0 mass matrix at position v70."""
  tf_scalar_evaluator = scalar_sector_tensorflow.get_tf_scalar_evaluator()
  t_left_onb = tf.constant(numpy.zeros(70), dtype=tf.float64)
  t_v70 = tf.constant(v70, dtype=tf.complex128)
  sinfo0 = tf_scalar_evaluator(t_v70)
  tc_v70_from_v70o = tf.constant(algebra.e7.v70_from_v70o)
  def tf_potential_relative(t_left_onb):
    sinfo = tf_scalar_evaluator(
      t_v70,
      t_left=tf.cast(tf.einsum('vV,V->v', tc_v70_from_v70o, t_left_onb),
                     tf.complex128))
    return sinfo.potential
  scalar_mm = m_util.tf_hessian(tf_potential_relative)(
    tf.constant(numpy.zeros(70), dtype=tf.float64)).numpy()
  return (scalar_mm /
          # Factor -3/8 (rather than -3/4) is due to normalization of
          # our orthonormal basis.
          (-3.0 / 8) / sinfo0.potential.numpy())
//  [python]
def string_to_int_tensor(x):
  """Graph operations decode a string into a fixed-size tensor of ints."""
  decoded = tf.compat.v1.decode_raw(x, tf.uint8)
  clipped = decoded[:STRING_TO_INT_MAX_CONTENTS_LEN]  # clip to allowed max_len
  shape = tf.shape(input=clipped)
  slen = shape[0]
  # pad to desired max_len
  padded = tf.pad(
      tensor=clipped, paddings=[[0, STRING_TO_INT_MAX_CONTENTS_LEN - slen]]
  )
  casted = tf.compat.v1.to_int32(padded)
  casted.set_shape([STRING_TO_INT_MAX_CONTENTS_LEN])
  return tf.concat([[slen], casted], 0)
//  [python]
def string_to_int_tensor(x):
  """Graph operations decode a string into a fixed-size tensor of ints."""
  decoded = tf.compat.v1.decode_raw(x, tf.uint8)
  clipped = decoded[:STRING_TO_INT_MAX_CONTENTS_LEN]  # clip to allowed max_len
  shape = tf.shape(input=clipped)
  slen = shape[0]
  # pad to desired max_len
  padded = tf.pad(
      tensor=clipped, paddings=[[0, STRING_TO_INT_MAX_CONTENTS_LEN - slen]]
  )
  casted = tf.cast(padded, tf.int32)
  casted.set_shape([STRING_TO_INT_MAX_CONTENTS_LEN])
  return tf.concat([[slen], casted], 0)
//  [python]
def fancy_slice_2d(X, inds0, inds1):
    """
    like numpy X[inds0, inds1]
    XXX this implementation is bad
    """
    inds0 = t.cast(inds0, t.int64)
    inds1 = t.cast(inds1, t.int64)
    shape = t.cast(t.shape(X), t.int64)
    ncols = shape[1]
    Xflat = t.reshape(X, [-1])
    return t.gather(Xflat, inds0 * ncols + inds1)
//  [python]
def fancy_slice_2d(X, inds0, inds1):
    """
    like numpy X[inds0, inds1]
    XXX this implementation is bad
    """
    inds0 = tf.cast(inds0, tf.int64)
    inds1 = tf.cast(inds1, tf.int64)
    shape = tf.cast(tf.shape(X), tf.int64)
    ncols = shape[1]
    Xflat = tf.reshape(X, [-1])
    return tf.gather(Xflat, inds0 * ncols + inds1)
//  [python]
def preprocess_images(images):
  """Applies preprocessing operations for Inception images.

  Because this will run in model_fn, on the accelerator, we use operations
  that efficiently execute there.

  Args:
    images: A Tensor of shape [batch_size height, width, channel] with uint8
      values.

  Returns:
    A tensor of images of shape [batch_size height, width, channel]
    containing floating point values, with all points rescaled between
    -1 and 1 and possibly resized.
  """
  images = tf.cast(images, dtype=tf.float32)
  images = tf.subtract(images, 128.0)
  images = tf.math.divide(images, 128.0)
  return images
//  [python]
def call_with_scalar_mass_matrix_evaluator(f, *args):
  """Returns f(evaluator, *args), with `evaluator` a mass matrix evaluator.

  Here, `evaluator` is a TensorFlow-based evaluator function
  that maps a 70-vector to a scalar mass matrix.

  We are then passing this on to a function that gets run in TensorFlow
  session context. This way, we can bulk-process solutions.
  """
  graph = tf.Graph()
  with graph.as_default():
    tf_scalar_evaluator = scalar_sector_tensorflow.get_tf_scalar_evaluator()
    t_left_onb = tf.Variable(  # Uses orthonormal basis.
        initial_value=numpy.zeros([70]), trainable=False, dtype=tf.float64)
    t_input = tf.placeholder(tf.float64, shape=[70])
    t_v70 = tf.Variable(
        initial_value=numpy.zeros([70]), trainable=False, dtype=tf.float64)
    op_assign_input = tf.assign(t_v70, t_input)
    sinfo = tf_scalar_evaluator(
        t.cast(t_v70, tf.complex128),
        t_left=t.cast(tf.einsum('vV,V->v',
                                 tf.constant(algebra.e7.v70_from_v70o),
                                 t_left_onb),
                       tf.complex128))
    t_potential = sinfo.potential
    t_scalar_mass_matrix = (
        tf.real(tf.hessians([t_potential], [t_left_onb])[0]) *
        # Factor -3/8 (rather than -3/4) is due to normalization of
        # our orthonormal basis.
        (-3.0 / 8) / t_potential)
    with tf.compat.v1.Session() as sess:
      sess.run([tf.global_variables_initializer()])
      def evaluator(v70):
        sess.run([op_assign_input], feed_dict={t_input: v70})
        ret = sess.run([t_scalar_mass_matrix])[0]
        return ret
      return f(evaluator, *args)
//  [python]
def call_with_scalar_mass_matrix_evaluator(f, *args):
  """Returns f(evaluator, *args), with `evaluator` a mass matrix evaluator.

  Here, `evaluator` is a TensorFlow-based evaluator function
  that maps a 70-vector to a scalar mass matrix.

  We are then passing this on to a function that gets run in TensorFlow
  session context. This way, we can bulk-process solutions.
  """
  graph = tf.Graph()
  with graph.as_default():
    tf_scalar_evaluator = scalar_sector_tensorflow.get_tf_scalar_evaluator()
    t_left_onb = tf.Variable(  # Uses orthonormal basis.
        initial_value=numpy.zeros([70]), trainable=False, dtype=tf.float64)
    t_input = tf.placeholder(tf.float64, shape=[70])
    t_v70 = tf.Variable(
        initial_value=numpy.zeros([70]), trainable=False, dtype=tf.float64)
    op_assign_input = tf.assign(t_v70, t_input)
    sinfo = tf_scalar_evaluator(
        tf.cast(t_v70, tf.complex128),
        t_left=tf.cast(tf.einsum('vV,V->v',
                                 tf.constant(algebra.e7.v70_from_v70o),
                                 t_left_onb),
                       tf.complex128))
    t_potential = sinfo.potential
    t_scalar_mass_matrix = (
        tf.real(tf.hessians([t_potential], [t_left_onb])[0]) *
        # Factor -3/8 (rather than -3/4) is due to normalization of
        # our orthonormal basis.
        (-3.0 / 8) / t_potential)
    with tf.compat.v1.Session() as sess:
      sess.run([tf.global_variables_initializer()])
      def evaluator(v70):
        sess.run([op_assign_input], feed_dict={t_input: v70})
        ret = sess.run([t_scalar_mass_matrix])[0]
        return ret
      return f(evaluator, *args)
//  [python]
def loss_layer(self, project_logits, lengths, name=None):
        """
        calculate crf loss
        :param project_logits: [1, num_steps, num_tags]
        :return: scalar loss
        """
        with tf.variable_scope("crf_loss"  if not name else name):
            small = -1000.0
            # pad logits for crf loss
            start_logits = tf.concat(
                [small * tf.ones(shape=[self.batch_size, 1, self.num_tags]), tf.zeros(shape=[self.batch_size, 1, 1])], axis=-1)
            pad_logits = tf.cast(small * tf.ones([self.batch_size, self.num_steps, 1]), tf.float32)
            logits = tf.concat([project_logits, pad_logits], axis=-1)
            logits = tf.concat([start_logits, logits], axis=1)
            targets = tf.concat(
                [tf.to_int32(self.num_tags*tf.ones([self.batch_size, 1])), self.targets], axis=-1)

            self.trans = tf.get_variable(
                "transitions",
                shape=[self.num_tags + 1, self.num_tags + 1],
                initializer=self.initializer)
            log_likelihood, self.trans = crf_log_likelihood(
                inputs=logits,
                tag_indices=targets,
                transition_params=self.trans,
                sequence_lengths=lengths+1)
            return tf.reduce_mean(-log_likelihood)
//  [python]
def loss_layer(self, project_logits, lengths, name=None):
        """
        calculate crf loss
        :param project_logits: [1, num_steps, num_tags]
        :return: scalar loss
        """
        with tf.variable_scope("crf_loss"  if not name else name):
            small = -1000.0
            # pad logits for crf loss
            start_logits = tf.concat(
                [small * tf.ones(shape=[self.batch_size, 1, self.num_tags]), tf.zeros(shape=[self.batch_size, 1, 1])], axis=-1)
            pad_logits = tf.cast(small * tf.ones([self.batch_size, self.num_steps, 1]), tf.float32)
            logits = tf.concat([project_logits, pad_logits], axis=-1)
            logits = tf.concat([start_logits, logits], axis=1)
            targets = tf.concat(
                [tf.cast(self.num_tags*tf.ones([self.batch_size, 1]), tf.int32), self.targets], axis=-1)

            self.trans = tf.get_variable(
                "transitions",
                shape=[self.num_tags + 1, self.num_tags + 1],
                initializer=self.initializer)
            log_likelihood, self.trans = crf_log_likelihood(
                inputs=logits,
                tag_indices=targets,
                transition_params=self.trans,
                sequence_lengths=lengths+1)
            return tf.reduce_mean(-log_likelihood)
//  [python]
def process_batch_input(input_batch):
  if FLAGS.dataset == 'celeba':
    return input_batch
  else:
    input_batch = tf.reshape(input_batch, [tf.shape(input_batch)[0], -1])
    input_batch = tf.to_float(input_batch)
    return input_batch
//  [python]
def process_batch_input(input_batch):
  if FLAGS.dataset == 'celeba':
    return input_batch
  else:
    input_batch = tf.reshape(input_batch, [tf.shape(input_batch)[0], -1])
    input_batch = tf.cast(input_batch, tf.float32)
    return input_batch
//  [python]
def call(
        self,
        hidden_states: tf.Tensor,
        head_mask: tf.Tensor,
        output_attentions: bool,
        training: bool = False,
    ) -> Tuple[tf.Tensor]:
        batch_size = shape_list(hidden_states)[0]
        mixed_query_layer = self.query(inputs=hidden_states)
        mixed_key_layer = self.key(inputs=hidden_states)
        mixed_value_layer = self.value(inputs=hidden_states)
        query_layer = self.transpose_for_scores(mixed_query_layer, batch_size)
        key_layer = self.transpose_for_scores(mixed_key_layer, batch_size)
        value_layer = self.transpose_for_scores(mixed_value_layer, batch_size)

        # Take the dot product between "query" and "key" to get the raw attention scores.
        # (batch size, num_heads, seq_len_q, seq_len_k)
        attention_scores = tf.matmul(query_layer, key_layer, transpose_b=True)
        dk = tf.to_float(self.sqrt_att_head_size)
        attention_scores = tf.divide(attention_scores, dk)

        # Normalize the attention scores to probabilities.
        attention_probs = tf.nn.softmax(logits=attention_scores, axis=-1)

        # This is actually dropping out entire tokens to attend to, which might
        # seem a bit unusual, but is taken from the original Transformer paper.
        attention_probs = self.dropout(inputs=attention_probs, training=training)

        # Mask heads if we want to
        if head_mask is not None:
            attention_probs = tf.multiply(attention_probs, head_mask)

        attention_output = tf.matmul(attention_probs, value_layer)
        attention_output = tf.transpose(attention_output, perm=[0, 2, 1, 3])

        # (batch_size, seq_len_q, all_head_size)
        attention_output = tf.reshape(tensor=attention_output, shape=(batch_size, -1, self.all_head_size))
        outputs = (attention_output, attention_probs) if output_attentions else (attention_output,)

        return outputs
//  [python]
def call(
        self,
        hidden_states: tf.Tensor,
        head_mask: tf.Tensor,
        output_attentions: bool,
        training: bool = False,
    ) -> Tuple[tf.Tensor]:
        batch_size = shape_list(hidden_states)[0]
        mixed_query_layer = self.query(inputs=hidden_states)
        mixed_key_layer = self.key(inputs=hidden_states)
        mixed_value_layer = self.value(inputs=hidden_states)
        query_layer = self.transpose_for_scores(mixed_query_layer, batch_size)
        key_layer = self.transpose_for_scores(mixed_key_layer, batch_size)
        value_layer = self.transpose_for_scores(mixed_value_layer, batch_size)

        # Take the dot product between "query" and "key" to get the raw attention scores.
        # (batch size, num_heads, seq_len_q, seq_len_k)
        attention_scores = tf.matmul(query_layer, key_layer, transpose_b=True)
        dk = tf.cast(self.sqrt_att_head_size, dtype=attention_scores.dtype)
        attention_scores = tf.divide(attention_scores, dk)

        # Normalize the attention scores to probabilities.
        attention_probs = tf.nn.softmax(logits=attention_scores, axis=-1)

        # This is actually dropping out entire tokens to attend to, which might
        # seem a bit unusual, but is taken from the original Transformer paper.
        attention_probs = self.dropout(inputs=attention_probs, training=training)

        # Mask heads if we want to
        if head_mask is not None:
            attention_probs = tf.multiply(attention_probs, head_mask)

        attention_output = tf.matmul(attention_probs, value_layer)
        attention_output = tf.transpose(attention_output, perm=[0, 2, 1, 3])

        # (batch_size, seq_len_q, all_head_size)
        attention_output = tf.reshape(tensor=attention_output, shape=(batch_size, -1, self.all_head_size))
        outputs = (attention_output, attention_probs) if output_attentions else (attention_output,)

        return outputs
//  [python]
def sinc(
    x: Union[tf.Tensor, tf.Variable],
    /,
    *,
    out: Optional[Union[tf.Tensor, tf.Variable]] = None,
) -> Union[tf.Tensor, tf.Variable]:
    x = ivy.pi * x
    return tf.where(x == 0, 1, tf.math.sin(x) / x).astype(x.dtype)
//  [python]
def sinc(
    x: Union[tf.Tensor, tf.Variable],
    /,
    *,
    out: Optional[Union[tf.Tensor, tf.Variable]] = None,
) -> Union[tf.Tensor, tf.Variable]:
    x = ivy.pi * x
    return tf.cast(tf.where(x == 0, 1, tf.math.sin(x) / x), x.dtype)
//  [python]
def infer_length(seq, pad_ix=0, time_major=False, dtype=tf.int32):
    """
    compute length given output indices and PAD code
    :param seq: tf matrix [time,batch] if time_major else [batch,time]
    :param pad_ix: integer index of end-of-sentence token
    :returns: lengths, int32 vector of shape [batch]
    """
    axis = 0 if time_major else 1
    is_pad = t.cast(t.equal(seq, pad_ix), dtype)
    count_pad = t.cumsum(is_pad, axis=axis, exclusive=False)
    lengths = t.reduce_sum(t.cast(t.equal(count_pad, 0), dtype), axis=axis)
    return lengths
//  [python]
def infer_length(seq, pad_ix=0, time_major=False, dtype=tf.int32):
    """
    compute length given output indices and PAD code
    :param seq: tf matrix [time,batch] if time_major else [batch,time]
    :param pad_ix: integer index of end-of-sentence token
    :returns: lengths, int32 vector of shape [batch]
    """
    axis = 0 if time_major else 1
    is_pad = tf.cast(tf.equal(seq, pad_ix), dtype)
    count_pad = tf.cumsum(is_pad, axis=axis, exclusive=False)
    lengths = tf.reduce_sum(tf.cast(tf.equal(count_pad, 0), dtype), axis=axis)
    return lengths
//  [python]
def __init__(self, layers, name=None):
       
        self.W = []  #Weights and biases
        self.parameters = 0 #total number of parameters
        
        for i in range(len(layers)-1):
            
            input_dim = layers[i]
            output_dim = layers[i+1]
            
            #Xavier standard deviation 
            std_dv = np.sqrt((2.0/(input_dim + output_dim)))

            #weights = normal distribution * Xavier standard deviation + 0
            w = tf.random.normal([input_dim, output_dim], dtype = 'float64') * std_dv
                        
            w = tf.Variable(w, trainable=True, name = 'w' + str(i+1))

            b = tf.Variable(tf.zeros([output_dim], dtype = 'float64'), trainable = True, name = 'b' + str(i+1))
                    
            self.W.append(w)
            self.W.append(b)
            
            self.parameters +=  input_dim * output_dim + output_dim
//  [python]
def __init__(self, layers, name=None):
       
        self.W = []  #Weights and biases
        self.parameters = 0 #total number of parameters
        
        for i in range(len(layers)-1):
            
            input_dim = layers[i]
            output_dim = layers[i+1]
            
            #Xavier standard deviation 
            std_dv = np.sqrt((2.0/(input_dim + output_dim)))

            #weights = normal distribution * Xavier standard deviation + 0
            w = tf.random.normal([input_dim, output_dim], dtype = 'float64') * std_dv
                        
            w = tf.Variable(w, trainable=True, name = 'w' + str(i+1))

            b = tf.Variable(tf.cast(tf.zeros([output_dim]), dtype = 'float64'), trainable = True, name = 'b' + str(i+1))
                    
            self.W.append(w)
            self.W.append(b)
            
            self.parameters +=  input_dim * output_dim + output_dim
//  [python]
def float_power(
    x1: Union[tf.Tensor, tf.Variable, float, list, tuple],
    x2: Union[tf.Tensor, tf.Variable, float, list, tuple],
    /,
    *,
    out: Optional[Union[tf.Tensor, tf.Variable]] = None,
) -> Union[tf.Tensor, tf.Variable]:
    x1, x2 = ivy.promote_types_of_inputs(x1, x2)
    if ivy.any(ivy.is_complex_dtype(x1)) or ivy.any(ivy.is_complex_dtype(x2)):
        out_dtype = tf.complex128
    else:
        out_dtype = tf.float64
    return tf.cast(tf.experimental.numpy.float_power(x1, x2), out_dtype)
//  [python]
def evaluate(model, tf_dataset, max_step=1000, num_eval_samples=None,
             stick_breaking=False,
             tree_structure=False):
  """Evaluate the model."""
  if num_eval_samples:
    num_samples = num_eval_samples
  elif FLAGS.num_eval_samples:
    num_samples = FLAGS.num_eval_samples
  else:
    num_samples = FLAGS.num_samples
  # tf.print('Evaluate with samples: %d.', num_samples)
  loss = 0.
  n = 0.
  for batch in tf_dataset.map(process_batch_input):
    if n >= max_step:  # used for train_ds, which is a `repeat` dataset.
      break
    if num_samples > 1:
      batch_size = tf.shape(batch)[0]
      input_batch = tf.tile(
          batch, [num_samples] + [1] * (len(batch.shape)-1))
      elbo = tf.reshape(model(input_batch,
                              stick_breaking=stick_breaking,
                              tree_structure=tree_structure)[0],
                        [num_samples, batch_size])
      objectives = (tf.reduce_logsumexp(elbo, axis=0, keepdims=False) -
                    tf.math.log(tf.cast(tf.shape(elbo)[0], tf.float32)))
    else:
      objectives = model(batch,
                         stick_breaking=stick_breaking,
                         tree_structure=tree_structure)[0]
    loss -= tf.reduce_mean(objectives)
    n += 1.
  return loss / n
//  [python]
def fix(
    x: Union[tf.Tensor, tf.Variable],
    /,
    *,
    out: Optional[Union[tf.Tensor, tf.Variable]] = None,
) -> Union[tf.Tensor, tf.Variable]:
    return tf.cast(tf.where(x > 0, tf.math.floor(x), tf.math.ceil(x)), x.dtype)
//  [python]
def copysign(
    x1: Union[tf.Tensor, tf.Variable, Number],
    x2: Union[tf.Tensor, tf.Variable, Number],
    /,
    *,
    out: Optional[tf.Tensor] = None,
) -> Union[tf.Tensor, tf.Variable]:
    x1, x2 = promote_types_of_inputs(x1, x2)
    # Cast our inputs to float to match numpy behaviour
    if not ivy.is_float_dtype(x1):
        x1 = t.cast(x1, ivy.default_float_dtype(as_native=True))
        x2 = t.cast(x2, ivy.default_float_dtype(as_native=True))
    # Replace any zero values with 1/the value, since tf.math.sign always
    # returns 0 for positive or negative zero
    signable_x2 = tf.where(tf.equal(x2, 0), tf.math.divide(1, x2), x2)
    signs = tf.math.sign(signable_x2)
    return tf.math.multiply(tf.math.abs(x1), signs)
//  [python]
def copysign(
    x1: Union[tf.Tensor, tf.Variable, Number],
    x2: Union[tf.Tensor, tf.Variable, Number],
    /,
    *,
    out: Optional[tf.Tensor] = None,
) -> Union[tf.Tensor, tf.Variable]:
    x1, x2 = promote_types_of_inputs(x1, x2)
    # Cast our inputs to float to match numpy behaviour
    if not ivy.is_float_dtype(x1):
        x1 = tf.cast(x1, ivy.default_float_dtype(as_native=True))
        x2 = tf.cast(x2, ivy.default_float_dtype(as_native=True))
    # Replace any zero values with 1/the value, since tf.math.sign always
    # returns 0 for positive or negative zero
    signable_x2 = tf.where(tf.equal(x2, 0), tf.math.divide(1, x2), x2)
    signs = tf.math.sign(signable_x2)
    return tf.math.multiply(tf.math.abs(x1), signs)
//  [python]
def train_one_step(
    train_batch_i,
    bvae_model,
    genmo_optimizer,
    infnet_optimizer,
    prior_optimizer,
    grad_variable_dict,
    grad_sq_variable_dict):
  """Train Discrete VAE for 1 step."""
  metrics = {}
  input_batch = process_batch_input(train_batch_i)

  (genmo_grads, prior_grads, infnet_grads, genmo_loss) = estimate_gradients(
      input_batch, bvae_model, FLAGS.grad_type,
      stick_breaking=FLAGS.stick_breaking,
      tree_structure=FLAGS.tree_structure,
      importance_weight=FLAGS.importance_weight,
      logits_sorting_order=FLAGS.logits_order)

  genmo_vars = bvae_model.decoder_vars
  genmo_optimizer.apply_gradients(list(zip(genmo_grads, genmo_vars)))

  prior_vars = bvae_model.prior_vars
  prior_optimizer.apply_gradients(list(zip(prior_grads, prior_vars)))

  infnet_vars = bvae_model.encoder_vars
  infnet_optimizer.apply_gradients(list(zip(infnet_grads, infnet_vars)))

  batch_size_sq = t.cast(FLAGS.batch_size * FLAGS.batch_size, t.float32)
  encoder_grad_var = bvae_model.compute_grad_variance(
      grad_variable_dict[FLAGS.grad_type],
      grad_sq_variable_dict[FLAGS.grad_type],
      infnet_grads) / batch_size_sq

  variance_dict = {}
  if (FLAGS.grad_type == 'reinforce_loo') and FLAGS.estimate_grad_basket:
    for method in FLAGS.estimate_grad_basket:
      if method == FLAGS.grad_type:
        continue  # Already computed
      if ('disarm' in method) and ('tree' not in method):
        main_method, logits_order = method.split('-')
        logits_order = None if logits_order == 'null' else logits_order
      (_, _, infnet_grads, _) = estimate_gradients(
          input_batch, bvae_model, main_method,
          stick_breaking=True,
          tree_structure=False,
          logits_sorting_order=logits_order)
      variance_dict[method] = bvae_model.compute_grad_variance(
          grad_variable_dict[method],
          grad_sq_variable_dict[method],
          infnet_grads) / batch_size_sq

  return (encoder_grad_var, variance_dict, genmo_loss, metrics)
//  [python]
def train_one_step(
    train_batch_i,
    bvae_model,
    genmo_optimizer,
    infnet_optimizer,
    prior_optimizer,
    grad_variable_dict,
    grad_sq_variable_dict):
  """Train Discrete VAE for 1 step."""
  metrics = {}
  input_batch = process_batch_input(train_batch_i)

  (genmo_grads, prior_grads, infnet_grads, genmo_loss) = estimate_gradients(
      input_batch, bvae_model, FLAGS.grad_type,
      stick_breaking=FLAGS.stick_breaking,
      tree_structure=FLAGS.tree_structure,
      importance_weight=FLAGS.importance_weight,
      logits_sorting_order=FLAGS.logits_order)

  genmo_vars = bvae_model.decoder_vars
  genmo_optimizer.apply_gradients(list(zip(genmo_grads, genmo_vars)))

  prior_vars = bvae_model.prior_vars
  prior_optimizer.apply_gradients(list(zip(prior_grads, prior_vars)))

  infnet_vars = bvae_model.encoder_vars
  infnet_optimizer.apply_gradients(list(zip(infnet_grads, infnet_vars)))

  batch_size_sq = tf.cast(FLAGS.batch_size * FLAGS.batch_size, tf.float32)
  encoder_grad_var = bvae_model.compute_grad_variance(
      grad_variable_dict[FLAGS.grad_type],
      grad_sq_variable_dict[FLAGS.grad_type],
      infnet_grads) / batch_size_sq

  variance_dict = {}
  if (FLAGS.grad_type == 'reinforce_loo') and FLAGS.estimate_grad_basket:
    for method in FLAGS.estimate_grad_basket:
      if method == FLAGS.grad_type:
        continue  # Already computed
      if ('disarm' in method) and ('tree' not in method):
        main_method, logits_order = method.split('-')
        logits_order = None if logits_order == 'null' else logits_order
      (_, _, infnet_grads, _) = estimate_gradients(
          input_batch, bvae_model, main_method,
          stick_breaking=True,
          tree_structure=False,
          logits_sorting_order=logits_order)
      variance_dict[method] = bvae_model.compute_grad_variance(
          grad_variable_dict[method],
          grad_sq_variable_dict[method],
          infnet_grads) / batch_size_sq

  return (encoder_grad_var, variance_dict, genmo_loss, metrics)
//  [python]
def extract_ho_conds_by_cls_index(self, cls_index, v_gt_class_HO):
        tmp_ho_index = tf.reduce_sum(tf.multiply(v_gt_class_HO,
                                                 tf.expand_dims(tf.to_float(cls_index), axis=0)), axis=-1)
        conds = tmp_ho_index > 0.0
        return conds
//  [python]
def extract_ho_conds_by_cls_index(self, cls_index, v_gt_class_HO):
        tmp_ho_index = tf.reduce_sum(tf.multiply(v_gt_class_HO,
                                                 tf.expand_dims(tf.cast(cls_index, tf.float32), axis=0)), axis=-1)
        conds = tmp_ho_index > 0.0
        return conds