jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def _graph_coarsening(self, X, A, layer, reuse):
        """Graph Coarsening used to reduce the graph scale

        Args:
            X (obj): node embedding in graph
            A (obj): adjacency matrix of graph
            layer (obj): graph coarsening layer
            reuse (obj): reusing attention W in query operation 

        Returns:
            X (obj): aggerated cluster embedding 
            A (obj): pooled adjacency matrix 
        """

        with tf.name_scope('CGAT'):
            # add self-loop
            A_bool = tf.greater(A, 0)
            A_bool = tf.cast(A_bool, A.dtype)
            A = A_bool * (tf.ones([self.max_n_nodes,self.max_n_nodes]) - tf.eye(self.max_n_nodes)) + tf.eye(self.max_n_nodes)

            # normalize graph
            D = tf.reduce_sum(A, axis=-1) # B*L
            D = tf.sqrt(D)[:, None] + K.epsilon() # B*1*L
            A = (A / D) / tf.transpose(D, perm=(0,2,1)) # B*L*L / B*1*L / B*L*1

            # attention propogation
            X_q = tf.matmul(A, X) # B*L*F
            X_q = tf.layers.dense(X_q, self.hparams.hidden_size)
            X_k = tf.layers.dense(X, self.hparams.hidden_size)
            f_1 = tf.layers.dense(X_q, 1)  # B*L*F x F*1 -> B*L*1
            f_2 = tf.layers.dense(X_k, 1) # B*L*F x F*1 -> B*L*1
            S = A_bool * f_1 + A_bool * tf.transpose(f_2, (0,2,1)) # B*L*1 x B*L*1 -> B*L*L
            S = tf.nn.leaky_relu(S)
            S = tf.nn.softmax(S, axis=-1) * A_bool
            X = tf.matmul(S, X) # B*L*L x B*L*F -> B*L*F

        with tf.name_scope('QPool'):
            if not self.query_shared:
                _, cluster_score = self._attention_fcn(self.target_item_embedding, X, layer, False, return_alpha=True)
            if self.query_shared:
                _, cluster_score = self._attention_fcn(self.target_item_embedding, X, 'shared', reuse, return_alpha=True)
            num_nodes = tf.reduce_sum(self.mask, 1) # B
            n_graphs = tf.shape(X)[0] 
            to_keep = tf.math.ceil(self.pool_ratio * tf.cast(num_nodes, tf.float32)) # B

            ## update mask
            cluster_score = cluster_score * self.real_mask # B*L
            if 'kwai' in socket.gethostname():
                sorted_score = tf.contrib.framework.sort(cluster_score, direction='DESCENDING', axis=-1) # B*L
            else:
                sorted_score = tf.sort(cluster_score, direction='DESCENDING', axis=-1) # B*L
            target_score = tf.gather_nd(sorted_score, tf.expand_dims(tf.cast(to_keep, tf.int32), -1), batch_dims=1)
            topk_mask = tf.greater_equal(cluster_score, tf.expand_dims(target_score, -1)) # B*L + B*1 -> B*L
            self.mask = tf.cast(topk_mask, tf.int32)
            self.real_mask = tf.cast(self.mask, tf.float32)

            ## ensure graph connectivity 
            S = S * tf.expand_dims(self.real_mask, -1) * tf.expand_dims(self.real_mask, -2)
            A = tf.matmul(tf.matmul(S, A),
                          tf.transpose(S, (0,2,1))) # B*C*L x B*L*L x B*L*C = B*C*C

        return X, A
//  [python]
def _graph_coarsening(self, X, A, layer, reuse):
        """Graph Coarsening used to reduce the graph scale

        Args:
            X (obj): node embedding in graph
            A (obj): adjacency matrix of graph
            layer (obj): graph coarsening layer
            reuse (obj): reusing attention W in query operation 

        Returns:
            X (obj): aggerated cluster embedding 
            A (obj): pooled adjacency matrix 
        """

        with tf.name_scope('CGAT'):
            # add self-loop
            A_bool = tf.cast(tf.greater(A, 0), A.dtype)
            A = A_bool * (tf.ones([self.max_n_nodes,self.max_n_nodes]) - tf.eye(self.max_n_nodes)) + tf.eye(self.max_n_nodes)

            # normalize graph
            D = tf.reduce_sum(A, axis=-1) # B*L
            D = tf.sqrt(D)[:, None] + K.epsilon() # B*1*L
            A = (A / D) / tf.transpose(D, perm=(0,2,1)) # B*L*L / B*1*L / B*L*1

            # attention propogation
            X_q = tf.matmul(A, X) # B*L*F
            X_q = tf.layers.dense(X_q, self.hparams.hidden_size)
            X_k = tf.layers.dense(X, self.hparams.hidden_size)
            f_1 = tf.layers.dense(X_q, 1)  # B*L*F x F*1 -> B*L*1
            f_2 = tf.layers.dense(X_k, 1) # B*L*F x F*1 -> B*L*1
            S = A_bool * f_1 + A_bool * tf.transpose(f_2, (0,2,1)) # B*L*1 x B*L*1 -> B*L*L
            S = tf.nn.leaky_relu(S)
            S = tf.nn.softmax(S, axis=-1) * A_bool
            X = tf.matmul(S, X) # B*L*L x B*L*F -> B*L*F

        with tf.name_scope('QPool'):
            if not self.query_shared:
                _, cluster_score = self._attention_fcn(self.target_item_embedding, X, layer, False, return_alpha=True)
            if self.query_shared:
                _, cluster_score = self._attention_fcn(self.target_item_embedding, X, 'shared', reuse, return_alpha=True)
            num_nodes = tf.reduce_sum(self.mask, 1) # B
            n_graphs = tf.shape(X)[0] 
            to_keep = tf.math.ceil(self.pool_ratio * tf.cast(num_nodes, tf.float32)) # B

            ## update mask
            cluster_score = cluster_score * self.real_mask # B*L
            if 'kwai' in socket.gethostname():
                sorted_score = tf.contrib.framework.sort(cluster_score, direction='DESCENDING', axis=-1) # B*L
            else:
                sorted_score = tf.sort(cluster_score, direction='DESCENDING', axis=-1) # B*L
            target_score = tf.gather_nd(sorted_score, tf.expand_dims(tf.cast(to_keep, tf.int32), -1), batch_dims=1)
            topk_mask = tf.greater_equal(cluster_score, tf.expand_dims(target_score, -1)) # B*L + B*1 -> B*L
            self.mask = tf.cast(topk_mask, tf.int32)
            self.real_mask = tf.cast(self.mask, tf.float32)

            ## ensure graph connectivity 
            S = S * tf.expand_dims(self.real_mask, -1) * tf.expand_dims(self.real_mask, -2)
            A = tf.matmul(tf.matmul(S, A),
                          tf.transpose(S, (0,2,1))) # B*C*L x B*L*L x B*L*C = B*C*C

        return X, A
//  [python]
def testScaleToZScore(self, input_data, output_data, elementwise):

    def preprocessing_fn(inputs):
      x = inputs['x']
      x_cast = t.cast(x, tf.as_dtype(input_data.dtype))
      x_scaled = tft.scale_to_z_score(x_cast, elementwise=elementwise)
      self.assertEqual(x_scaled.dtype, tf.as_dtype(output_data.dtype))
      return {'x_scaled': tf.cast(x_scaled, tf.float32)}

    input_data_dicts = [{'x': x} for x in input_data]
    expected_data_dicts = [{'x_scaled': x_scaled} for x_scaled in output_data]
    input_metadata = tft.DatasetMetadata.from_feature_spec({
        'x':
            tf.io.FixedLenFeature(
                input_data.shape[1:],
                tft_unit.canonical_numeric_dtype(tf.as_dtype(
                    input_data.dtype))),
    })
    expected_metadata = tft.DatasetMetadata.from_feature_spec({
        'x_scaled': tf.io.FixedLenFeature(output_data.shape[1:], tf.float32),
    })
    self.assertAnalyzeAndTransformResults(
        input_data_dicts, input_metadata,
        preprocessing_fn, expected_data_dicts, expected_metadata)
//  [python]
def testScaleToZScore(self, input_data, output_data, elementwise):

    def preprocessing_fn(inputs):
      x = inputs['x']
      x_cast = tf.cast(x, tf.as_dtype(input_data.dtype))
      x_scaled = tft.scale_to_z_score(x_cast, elementwise=elementwise)
      self.assertEqual(x_scaled.dtype, tf.as_dtype(output_data.dtype))
      return {'x_scaled': tf.cast(x_scaled, tf.float32)}

    input_data_dicts = [{'x': x} for x in input_data]
    expected_data_dicts = [{'x_scaled': x_scaled} for x_scaled in output_data]
    input_metadata = tft.DatasetMetadata.from_feature_spec({
        'x':
            tf.io.FixedLenFeature(
                input_data.shape[1:],
                tft_unit.canonical_numeric_dtype(tf.as_dtype(
                    input_data.dtype))),
    })
    expected_metadata = tft.DatasetMetadata.from_feature_spec({
        'x_scaled': tf.io.FixedLenFeature(output_data.shape[1:], tf.float32),
    })
    self.assertAnalyzeAndTransformResults(
        input_data_dicts, input_metadata,
        preprocessing_fn, expected_data_dicts, expected_metadata)
//  [python]
def testScaleToZScorePerKeyElementwise(self, input_data, expected_data,
                                         input_metadata, expected_metadata):

    def preprocessing_fn(inputs):
      outputs = {}
      outputs['x_scaled'] = tft.scale_to_z_score_per_key(
          t.cast(inputs['x'], tf.float32),
          key=inputs['key'],
          elementwise=True,
          key_vocabulary_filename=None)
      self.assertEqual(outputs['x_scaled'].dtype, tf.float32)
      return outputs

    self.assertAnalyzeAndTransformResults(input_data, input_metadata,
                                          preprocessing_fn, expected_data,
                                          expected_metadata)
//  [python]
def testScaleToZScorePerKeyElementwise(self, input_data, expected_data,
                                         input_metadata, expected_metadata):

    def preprocessing_fn(inputs):
      outputs = {}
      outputs['x_scaled'] = tft.scale_to_z_score_per_key(
          tf.cast(inputs['x'], tf.float32),
          key=inputs['key'],
          elementwise=True,
          key_vocabulary_filename=None)
      self.assertEqual(outputs['x_scaled'].dtype, tf.float32)
      return outputs

    self.assertAnalyzeAndTransformResults(input_data, input_metadata,
                                          preprocessing_fn, expected_data,
                                          expected_metadata)
//  [python]
def testScalePerKeyRagged(self, input_data, make_key_spec, expected_data,
                            input_dtype):
    make_x_spec = lambda: tf.io.RaggedFeature(  # pylint: disable=g-long-lambda
        tft_unit.canonical_numeric_dtype(input_dtype),
        value_key='x_val',
        partitions=[
            tf.io.RaggedFeature.RowLengths('x_row_lengths')  # pytype: disable=attribute-error
        ])
    input_metadata = tft.DatasetMetadata.from_feature_spec({
        'x': tft_unit.make_feature_spec_wrapper(make_x_spec),
        'key': tft_unit.make_feature_spec_wrapper(make_key_spec)
    })

    def preprocessing_fn(inputs):
      scaled_to_z_score = tft.scale_to_z_score_per_key(
          t.cast(inputs['x'], input_dtype), inputs['key'])
      self.assertEqual(scaled_to_z_score.dtype, _mean_output_dtype(input_dtype))
      return {
          'scaled_by_min_max':
              tft.scale_by_min_max_per_key(
                  t.cast(inputs['x'], input_dtype),
                  inputs['key'],
                  output_min=-1,
                  output_max=1),
          'scaled_to_0_1':
              tft.scale_to_0_1_per_key(
                  t.cast(inputs['x'], input_dtype), inputs['key']),
          'scaled_to_z_score':
              t.cast(scaled_to_z_score, tf.float32),
      }

    expected_specs = {}
    for output_name in ('scaled_by_min_max', 'scaled_to_0_1',
                        'scaled_to_z_score'):
      expected_specs[output_name] = tf.io.RaggedFeature(
          tf.float32,
          value_key='{}$ragged_values'.format(output_name),
          partitions=[
              tf.io.RaggedFeature.RowLengths(  # pytype: disable=attribute-error
                  '{}$row_lengths_1'.format(output_name))
          ])
    expected_metadata = tft.DatasetMetadata.from_feature_spec(expected_specs)
    self.assertAnalyzeAndTransformResults(input_data, input_metadata,
                                          preprocessing_fn, expected_data,
                                          expected_metadata)
//  [python]
def testScalePerKeyRagged(self, input_data, make_key_spec, expected_data,
                            input_dtype):
    make_x_spec = lambda: tf.io.RaggedFeature(  # pylint: disable=g-long-lambda
        tft_unit.canonical_numeric_dtype(input_dtype),
        value_key='x_val',
        partitions=[
            tf.io.RaggedFeature.RowLengths('x_row_lengths')  # pytype: disable=attribute-error
        ])
    input_metadata = tft.DatasetMetadata.from_feature_spec({
        'x': tft_unit.make_feature_spec_wrapper(make_x_spec),
        'key': tft_unit.make_feature_spec_wrapper(make_key_spec)
    })

    def preprocessing_fn(inputs):
      scaled_to_z_score = tft.scale_to_z_score_per_key(
          tf.cast(inputs['x'], input_dtype), inputs['key'])
      self.assertEqual(scaled_to_z_score.dtype, _mean_output_dtype(input_dtype))
      return {
          'scaled_by_min_max':
              tft.scale_by_min_max_per_key(
                  tf.cast(inputs['x'], input_dtype),
                  inputs['key'],
                  output_min=-1,
                  output_max=1),
          'scaled_to_0_1':
              tft.scale_to_0_1_per_key(
                  tf.cast(inputs['x'], input_dtype), inputs['key']),
          'scaled_to_z_score':
              tf.cast(scaled_to_z_score, tf.float32),
      }

    expected_specs = {}
    for output_name in ('scaled_by_min_max', 'scaled_to_0_1',
                        'scaled_to_z_score'):
      expected_specs[output_name] = tf.io.RaggedFeature(
          tf.float32,
          value_key='{}$ragged_values'.format(output_name),
          partitions=[
              tf.io.RaggedFeature.RowLengths(  # pytype: disable=attribute-error
                  '{}$row_lengths_1'.format(output_name))
          ])
    expected_metadata = tft.DatasetMetadata.from_feature_spec(expected_specs)
    self.assertAnalyzeAndTransformResults(input_data, input_metadata,
                                          preprocessing_fn, expected_data,
                                          expected_metadata)
//  [python]
def testNumericMeanWithSparseTensorReduceFalseOverflow(self):
    self._SkipIfOutputRecordBatches()

    def analyzer_fn(inputs):
      return {'mean': tft.mean(t.cast(inputs['sparse'], tf.int32), False)}

    input_data = [
        {'idx': [0, 1], 'val': [1, 1]},
        {'idx': [1, 3], 'val': [2147483647, 3]},
    ]
    input_metadata = tft.DatasetMetadata.from_feature_spec(
        {'sparse': tf.io.SparseFeature('idx', 'val', tf.int64, 4)})
    expected_outputs = {
        'mean': np.array([1., 1073741824., float('nan'), 3.], np.float32)
    }
    self.assertAnalyzerOutputs(input_data, input_metadata, analyzer_fn,
                               expected_outputs)
//  [python]
def testNumericMeanWithSparseTensorReduceFalseOverflow(self):
    self._SkipIfOutputRecordBatches()

    def analyzer_fn(inputs):
      return {'mean': tft.mean(tf.cast(inputs['sparse'], tf.int32), False)}

    input_data = [
        {'idx': [0, 1], 'val': [1, 1]},
        {'idx': [1, 3], 'val': [2147483647, 3]},
    ]
    input_metadata = tft.DatasetMetadata.from_feature_spec(
        {'sparse': tf.io.SparseFeature('idx', 'val', tf.int64, 4)})
    expected_outputs = {
        'mean': np.array([1., 1073741824., float('nan'), 3.], np.float32)
    }
    self.assertAnalyzerOutputs(input_data, input_metadata, analyzer_fn,
                               expected_outputs)
//  [python]
def row_mean_rec_rank(y_true, y_pred):
  """The Mean Reciprocal Rank metric.

  Args:
    y_true: Ground truth class labels with shape=(None,) and dtype=tf.int64.
    y_pred: Predicted class logits or probabilities with shape=(None, n_classes)
      and dtype=tf.float32.

  Returns:
    An array containing the reciprocal rank of the true class for each test
    instance.
  """
  assert tf.rank(y_true) == 1 and tf.rank(y_pred) == 2
  ranked_classes = tf.to_int64(
      tf.argsort(y_pred, axis=1, direction='DESCENDING'))
  mean_rec_rank = 1 / (tf.where(ranked_classes == y_true[:, None])[:, 1] + 1)
  return mean_rec_rank
//  [python]
def row_mean_rec_rank(y_true, y_pred):
  """The Mean Reciprocal Rank metric.

  Args:
    y_true: Ground truth class labels with shape=(None,) and dtype=tf.int64.
    y_pred: Predicted class logits or probabilities with shape=(None, n_classes)
      and dtype=tf.float32.

  Returns:
    An array containing the reciprocal rank of the true class for each test
    instance.
  """
  assert tf.rank(y_true) == 1 and tf.rank(y_pred) == 2
  ranked_classes = tf.cast(
      tf.argsort(y_pred, axis=1, direction='DESCENDING'), dtype=tf.int64)
  mean_rec_rank = 1 / (tf.where(ranked_classes == y_true[:, None])[:, 1] + 1)
  return mean_rec_rank
//  [python]
def testScaleToZScoreSparsePerDenseKey(self, input_dtype):
    # TODO(b/131852830) Add elementwise tests.
    def preprocessing_fn(inputs):

      def scale_to_z_score_per_key(tensor, key):
        z_score = tft.scale_to_z_score_per_key(
            tf.cast(tensor, input_dtype), key=key, elementwise=False)
        self.assertEqual(z_score.dtype, _mean_output_dtype(input_dtype))
        return tf.cast(z_score, tf.float32)

      return {
          'x_scaled': scale_to_z_score_per_key(inputs['x'], inputs['key']),
          'y_scaled': scale_to_z_score_per_key(inputs['y'], inputs['key']),
      }
    np_dtype = input_dtype.as_numpy_dtype
    input_data = [{
        'x': np.array([-4, 2], dtype=np_dtype),
        'y': np.array([0, 0], dtype=np_dtype),
        'key': 'a',
    }, {
        'x': np.array([10, 4], dtype=np_dtype),
        'y': np.array([0, 0], dtype=np_dtype),
        'key': 'a',
    }, {
        'x': np.array([1, -1], dtype=np_dtype),
        'y': np.array([0, 0], dtype=np_dtype),
        'key': 'b',
    }]
    # Mean(x) = 3, Mean(y) = 0
    # Var(x) = (-7^2 + -1^2 + 7^2 + 1^2) / 4 = 25, Var(y) = 0
    # StdDev(x) = 5, StdDev(y) = 0
    # 'b':
    # Mean(x) = 0, Mean(y) = 0
    # Var(x) = 1, Var(y) = 0
    # StdDev(x) = 1, StdDev(y) = 0
    expected_data = [
        {
            'x_scaled': [-1.4, -.2],  # [(-4 - 3) / 5, (2 - 3) / 5]
            'y_scaled': [0., 0.],
        },
        {
            'x_scaled': [1.4, .2],  # [(10 - 3) / 5, (4 - 3) / 5]
            'y_scaled': [0., 0.],
        },
        {
            'x_scaled': [1., -1.],  # [(1 - 0) / 1, (-1 - 0) / 1]
            'y_scaled': [0., 0.],
        }
    ]

    input_metadata = tft.DatasetMetadata.from_feature_spec({
        'x': tf.io.VarLenFeature(tft_unit.canonical_numeric_dtype(input_dtype)),
        'y': tf.io.VarLenFeature(tft_unit.canonical_numeric_dtype(input_dtype)),
        'key': tf.io.FixedLenFeature([], tf.string),
    })
    expected_metadata = tft.DatasetMetadata.from_feature_spec({
        'x_scaled': tf.io.VarLenFeature(tf.float32),
        'y_scaled': tf.io.VarLenFeature(tf.float32),
    })
    self.assertAnalyzeAndTransformResults(input_data, input_metadata,
                                          preprocessing_fn, expected_data,
                                          expected_metadata)
//  [python]
def testSavedModelWithGlobalAnnotations(self):
    self._SkipIfOutputRecordBatches()

    def preprocessing_fn(inputs):
      # Add some arbitrary annotation data at the global schema level.
      boundaries = tf.constant([[1.0]])
      message_type = annotations_pb2.BucketBoundaries.DESCRIPTOR.full_name
      sizes = tf.expand_dims([tf.size(boundaries)], axis=0)
      message_proto = tf.raw_ops.EncodeProto(
          sizes=sizes, values=[t.float32(boundaries)],
          field_names=['boundaries'], message_type=message_type)[0]
      type_url = os.path.join('type.googleapis.com', message_type)
      schema_inference.annotate(type_url, message_proto)
      return {
          'x_scaled': tft.scale_by_min_max(inputs['x']),
      }

    input_data = [{'x': 1}, {'x': 2}, {'x': 3}, {'x': 4}]
    input_metadata = tft.DatasetMetadata.from_feature_spec({
        'x': tf.io.FixedLenFeature([], tf.float32),
    })
    temp_dir = self.get_temp_dir()
    with tft_beam.Context(temp_dir=temp_dir):
      transform_fn = ((input_data, input_metadata)
                      | tft_beam.AnalyzeDataset(preprocessing_fn))
      #  Write transform_fn to serialize annotation collections to SavedModel
      _ = transform_fn | transform_fn_io.WriteTransformFn(temp_dir)

    # Ensure that global annotations survive the round trip to SavedModel.
    tf_transform_output = tft.TFTransformOutput(temp_dir)
    schema = tf_transform_output.transformed_metadata._schema
    self.assertLen(schema.annotation.extra_metadata, 1)
    for annotation in schema.annotation.extra_metadata:
      message = annotations_pb2.BucketBoundaries()
      annotation.Unpack(message)
      self.assertAllClose(list(message.boundaries), [1])
//  [python]
def testSavedModelWithGlobalAnnotations(self):
    self._SkipIfOutputRecordBatches()

    def preprocessing_fn(inputs):
      # Add some arbitrary annotation data at the global schema level.
      boundaries = tf.constant([[1.0]])
      message_type = annotations_pb2.BucketBoundaries.DESCRIPTOR.full_name
      sizes = tf.expand_dims([tf.size(boundaries)], axis=0)
      message_proto = tf.raw_ops.EncodeProto(
          sizes=sizes, values=[tf.cast(boundaries, tf.float32)],
          field_names=['boundaries'], message_type=message_type)[0]
      type_url = os.path.join('type.googleapis.com', message_type)
      schema_inference.annotate(type_url, message_proto)
      return {
          'x_scaled': tft.scale_by_min_max(inputs['x']),
      }

    input_data = [{'x': 1}, {'x': 2}, {'x': 3}, {'x': 4}]
    input_metadata = tft.DatasetMetadata.from_feature_spec({
        'x': tf.io.FixedLenFeature([], tf.float32),
    })
    temp_dir = self.get_temp_dir()
    with tft_beam.Context(temp_dir=temp_dir):
      transform_fn = ((input_data, input_metadata)
                      | tft_beam.AnalyzeDataset(preprocessing_fn))
      #  Write transform_fn to serialize annotation collections to SavedModel
      _ = transform_fn | transform_fn_io.WriteTransformFn(temp_dir)

    # Ensure that global annotations survive the round trip to SavedModel.
    tf_transform_output = tft.TFTransformOutput(temp_dir)
    schema = tf_transform_output.transformed_metadata._schema
    self.assertLen(schema.annotation.extra_metadata, 1)
    for annotation in schema.annotation.extra_metadata:
      message = annotations_pb2.BucketBoundaries()
      annotation.Unpack(message)
      self.assertAllClose(list(message.boundaries), [1])
//  [python]
def testNumericAnalyzersWithScalarInputs(self, input_dtype, output_dtypes):
    self._SkipIfOutputRecordBatches()

    def analyzer_fn(inputs):
      a = t.cast(inputs['a'], input_dtype)

      def assert_and_cast_dtype(tensor, out_dtype):
        self.assertEqual(tensor.dtype, out_dtype)
        return t.cast(tensor, tft_unit.canonical_numeric_dtype(out_dtype))

      return {
          'min': assert_and_cast_dtype(tft.min(a),
                                       output_dtypes['min']),
          'max': assert_and_cast_dtype(tft.max(a),
                                       output_dtypes['max']),
          'sum': assert_and_cast_dtype(tft.sum(a),
                                       output_dtypes['sum']),
          'size': assert_and_cast_dtype(tft.size(a),
                                        output_dtypes['size']),
          'mean': assert_and_cast_dtype(tft.mean(a),
                                        output_dtypes['mean']),
          'var': assert_and_cast_dtype(tft.var(a),
                                       output_dtypes['var']),
      }

    input_data = [{'a': 4}, {'a': 1}]
    input_metadata = tft.DatasetMetadata.from_feature_spec({
        'a':
            tf.io.FixedLenFeature([],
                                  tft_unit.canonical_numeric_dtype(input_dtype))
    })
    expected_outputs = {
        'min':
            np.array(
                1,
                tft_unit.canonical_numeric_dtype(
                    output_dtypes['min']).as_numpy_dtype),
        'max':
            np.array(
                4,
                tft_unit.canonical_numeric_dtype(
                    output_dtypes['max']).as_numpy_dtype),
        'sum':
            np.array(
                5,
                tft_unit.canonical_numeric_dtype(
                    output_dtypes['sum']).as_numpy_dtype),
        'size':
            np.array(
                2,
                tft_unit.canonical_numeric_dtype(
                    output_dtypes['size']).as_numpy_dtype),
        'mean':
            np.array(
                2.5,
                tft_unit.canonical_numeric_dtype(
                    output_dtypes['mean']).as_numpy_dtype),
        'var':
            np.array(
                2.25,
                tft_unit.canonical_numeric_dtype(
                    output_dtypes['var']).as_numpy_dtype),
    }

    self.assertAnalyzerOutputs(
        input_data, input_metadata, analyzer_fn, expected_outputs)
//  [python]
def testNumericAnalyzersWithScalarInputs(self, input_dtype, output_dtypes):
    self._SkipIfOutputRecordBatches()

    def analyzer_fn(inputs):
      a = tf.cast(inputs['a'], input_dtype)

      def assert_and_cast_dtype(tensor, out_dtype):
        self.assertEqual(tensor.dtype, out_dtype)
        return tf.cast(tensor, tft_unit.canonical_numeric_dtype(out_dtype))

      return {
          'min': assert_and_cast_dtype(tft.min(a),
                                       output_dtypes['min']),
          'max': assert_and_cast_dtype(tft.max(a),
                                       output_dtypes['max']),
          'sum': assert_and_cast_dtype(tft.sum(a),
                                       output_dtypes['sum']),
          'size': assert_and_cast_dtype(tft.size(a),
                                        output_dtypes['size']),
          'mean': assert_and_cast_dtype(tft.mean(a),
                                        output_dtypes['mean']),
          'var': assert_and_cast_dtype(tft.var(a),
                                       output_dtypes['var']),
      }

    input_data = [{'a': 4}, {'a': 1}]
    input_metadata = tft.DatasetMetadata.from_feature_spec({
        'a':
            tf.io.FixedLenFeature([],
                                  tft_unit.canonical_numeric_dtype(input_dtype))
    })
    expected_outputs = {
        'min':
            np.array(
                1,
                tft_unit.canonical_numeric_dtype(
                    output_dtypes['min']).as_numpy_dtype),
        'max':
            np.array(
                4,
                tft_unit.canonical_numeric_dtype(
                    output_dtypes['max']).as_numpy_dtype),
        'sum':
            np.array(
                5,
                tft_unit.canonical_numeric_dtype(
                    output_dtypes['sum']).as_numpy_dtype),
        'size':
            np.array(
                2,
                tft_unit.canonical_numeric_dtype(
                    output_dtypes['size']).as_numpy_dtype),
        'mean':
            np.array(
                2.5,
                tft_unit.canonical_numeric_dtype(
                    output_dtypes['mean']).as_numpy_dtype),
        'var':
            np.array(
                2.25,
                tft_unit.canonical_numeric_dtype(
                    output_dtypes['var']).as_numpy_dtype),
    }

    self.assertAnalyzerOutputs(
        input_data, input_metadata, analyzer_fn, expected_outputs)
//  [python]
def label_confidence_change(
    intervened_ds, concept_name_to_intervene,
    curr_intervention_mask, signed):
  """Metric that measures the importance of a concept using the expected change in the predicted label confidence.

  Args:
    intervened_ds: The IntervenedDataset object to use for metric computation.
    concept_name_to_intervene: The name of the concept group to intervene on in
      the next intervention step.
    curr_intervention_mask: Intervention mask being used for the current
      intervention step.
    signed: Whether to use the signed difference for entropy to measure change.

  Returns:
    The expected change in prediced label confidence.
  """
  concept_probs, dist, dist_0, dist_1 = label_metric_utils(
      intervened_ds, concept_name_to_intervene, curr_intervention_mask)

  if intervened_ds.n_classes == 1:
    pred_class = t.cast(dist.logits > 0, t.float32)
  else:
    pred_class = tf.argmax(dist.logits, axis=1)
  pred_class_prob = dist.prob(pred_class)
  pred_class_prob_0 = dist_0.prob(pred_class)
  pred_class_prob_1 = dist_1.prob(pred_class)
  change_0 = pred_class_prob_0 - pred_class_prob
  change_1 = pred_class_prob_1 - pred_class_prob

  if signed:
    metric = concept_probs * change_1 + (1 - concept_probs) * change_0
  else:
    metric = (concept_probs * tf.abs(change_1)
              + (1 - concept_probs) * tf.abs(change_0))

  return metric
//  [python]
def label_confidence_change(
    intervened_ds, concept_name_to_intervene,
    curr_intervention_mask, signed):
  """Metric that measures the importance of a concept using the expected change in the predicted label confidence.

  Args:
    intervened_ds: The IntervenedDataset object to use for metric computation.
    concept_name_to_intervene: The name of the concept group to intervene on in
      the next intervention step.
    curr_intervention_mask: Intervention mask being used for the current
      intervention step.
    signed: Whether to use the signed difference for entropy to measure change.

  Returns:
    The expected change in prediced label confidence.
  """
  concept_probs, dist, dist_0, dist_1 = label_metric_utils(
      intervened_ds, concept_name_to_intervene, curr_intervention_mask)

  if intervened_ds.n_classes == 1:
    pred_class = tf.cast(dist.logits > 0, tf.float32)
  else:
    pred_class = tf.argmax(dist.logits, axis=1)
  pred_class_prob = dist.prob(pred_class)
  pred_class_prob_0 = dist_0.prob(pred_class)
  pred_class_prob_1 = dist_1.prob(pred_class)
  change_0 = pred_class_prob_0 - pred_class_prob
  change_1 = pred_class_prob_1 - pred_class_prob

  if signed:
    metric = concept_probs * change_1 + (1 - concept_probs) * change_0
  else:
    metric = (concept_probs * tf.abs(change_1)
              + (1 - concept_probs) * tf.abs(change_0))

  return metric
//  [python]
def test_ordering_rv_tensor(self):
    # Check that random variables are copied correctly in dependency
    # structure.
    with self.test_session() as sess:
      ed.set_seed(12432)
      x = Bernoulli(logits=0.0)
      y = t.cast(x, tf.float32)
      y_new = ed.copy(y)
      x_new = ed.copy(x)
      x_new_val, y_new_val = sess.run([x_new, y_new])
      self.assertEqual(x_new_val, y_new_val)
//  [python]
def test_ordering_rv_tensor(self):
    # Check that random variables are copied correctly in dependency
    # structure.
    with self.test_session() as sess:
      ed.set_seed(12432)
      x = Bernoulli(logits=0.0)
      y = tf.cast(x, tf.float32)
      y_new = ed.copy(y)
      x_new = ed.copy(x)
      x_new_val, y_new_val = sess.run([x_new, y_new])
      self.assertEqual(x_new_val, y_new_val)
//  [python]
def call(self, inputs, outer_rank, training=False, mask=None):
    # outer_rank is needed because the projection is not done on the raw
    # observations so getting the outer rank is hard as there is no spec to
    # compare to.
    batch_squash = utils.BatchSquash(outer_rank)
    inputs = batch_squash.flatten(inputs)
    inputs = t.cast(inputs, t.float32)

    logits = self._projection_layer(inputs, training=training)
    logits = tf.reshape(logits, [-1] + self._output_shape.as_list())
    logits = batch_squash.unflatten(logits)

    if mask is not None:
      # If the action spec says each action should be shaped (1,), add another
      # dimension so the final shape is (B, 1, A), where A is the number of
      # actions. This will make Categorical emit events shaped (B, 1) rather
      # than (B,). Using axis -2 to allow for (B, T, 1, A) shaped q_values.
      if mask.shape.rank < logits.shape.rank:
        mask = tf.expand_dims(mask, -2)

      # Overwrite the logits for invalid actions to a very large negative
      # number. We do not use -inf because it produces NaNs in many tfp
      # functions.
      almost_neg_inf = tf.constant(logits.dtype.min, dtype=logits.dtype)
      logits = tf.compat.v2.where(
          tf.cast(mask, tf.bool), logits, almost_neg_inf
      )

    return self.output_spec.build_distribution(logits=logits), ()
//  [python]
def call(self, inputs, outer_rank, training=False, mask=None):
    # outer_rank is needed because the projection is not done on the raw
    # observations so getting the outer rank is hard as there is no spec to
    # compare to.
    batch_squash = utils.BatchSquash(outer_rank)
    inputs = batch_squash.flatten(inputs)
    inputs = tf.cast(inputs, tf.float32)

    logits = self._projection_layer(inputs, training=training)
    logits = tf.reshape(logits, [-1] + self._output_shape.as_list())
    logits = batch_squash.unflatten(logits)

    if mask is not None:
      # If the action spec says each action should be shaped (1,), add another
      # dimension so the final shape is (B, 1, A), where A is the number of
      # actions. This will make Categorical emit events shaped (B, 1) rather
      # than (B,). Using axis -2 to allow for (B, T, 1, A) shaped q_values.
      if mask.shape.rank < logits.shape.rank:
        mask = tf.expand_dims(mask, -2)

      # Overwrite the logits for invalid actions to a very large negative
      # number. We do not use -inf because it produces NaNs in many tfp
      # functions.
      almost_neg_inf = tf.constant(logits.dtype.min, dtype=logits.dtype)
      logits = tf.compat.v2.where(
          tf.cast(mask, tf.bool), logits, almost_neg_inf
      )

    return self.output_spec.build_distribution(logits=logits), ()
//  [python]
def label_metric_utils(
    # Using Any for type annotation instead of
    # intervened_dataset.IntervenedDataset to avoid circular imports.
    intervened_ds, concept_name_to_intervene,
    curr_intervention_mask
):
  """Common utilities for computing concept importance metrics.



  Args:
    intervened_ds: The IntervenedDataset object to use for metric computation.
    concept_name_to_intervene: The name of the concept group to intervene on in
      the next intervention step.
    curr_intervention_mask: Intervention mask being used for the current
      intervention step.

  Returns:
    concept_probs: The mean predicted probability of the concept group, for each
      test example.
    dist: The predicted class label distribution using <curr_intervention_mask>.
    dist_0: The predicted class label distribution using an intervention mask
      that also reveals <concept_name_to_intervene> concept group assuiming 0 to
      be the true target for all binary concepts in the group.
    dist_1: The predicted class label distribution using an intervention mask
      that also reveals <concept_name_to_intervene> concept group assuiming 1 to
      be the true target for all binary concepts in the group.
  """
  b_concepts_to_reveal = intervened_ds.concept_groups[concept_name_to_intervene]
  concept_probs = tf.reduce_mean(
      tf.sigmoid(
          tf.gather(intervened_ds.pred_concepts, indices=b_concepts_to_reveal,
                    axis=1)),
      axis=1)
  b_concepts_to_reveal_multihot = tf.cast(
      tf.reduce_sum(
          tf.one_hot(b_concepts_to_reveal, intervened_ds.n_concepts), axis=0),
      dtype=tf.bool)
  mask_temp = tf.where(
      b_concepts_to_reveal_multihot, True, curr_intervention_mask)
  true_concepts_0 = tf.where(
      b_concepts_to_reveal_multihot, 0, intervened_ds.true_concepts)
  true_concepts_1 = tf.where(
      b_concepts_to_reveal_multihot, 1, intervened_ds.true_concepts)
  intervened_concepts = intervened_ds.intervene_on_batch(
      intervened_ds.true_concepts, intervened_ds.concept_uncertainty,
      intervened_ds.pred_concepts, curr_intervention_mask)
  intervened_concepts_0 = intervened_ds.intervene_on_batch(
      true_concepts_0, intervened_ds.concept_uncertainty,
      intervened_ds.pred_concepts, mask_temp)
  intervened_concepts_1 = intervened_ds.intervene_on_batch(
      true_concepts_1, intervened_ds.concept_uncertainty,
      intervened_ds.pred_concepts, mask_temp)
  if intervened_ds.n_classes == 1:
    dist_class = tfd.Bernoulli
  else:
    dist_class = tfd.Categorical
  dist = dist_class(
      logits=tf.squeeze(
          intervened_ds.ctoy_model(intervened_concepts, training=False)[0]))
  dist_0 = dist_class(
      logits=tf.squeeze(
          intervened_ds.ctoy_model(intervened_concepts_0, training=False)[0]))
  dist_1 = dist_class(
      logits=tf.squeeze(
          intervened_ds.ctoy_model(intervened_concepts_1, training=False)[0]))
  return concept_probs, dist, dist_0, dist_1
//  [python]
def label_confidence_changev2(
    intervened_ds, concept_name_to_intervene,
    curr_intervention_mask, signed):
  """Metric that measures the importance of a concept using the expected change in the predicted label confidence.

  This v2 method deals with categorical concepts differently by taking
  expectation under the distribution each predicted binary concept in the
  <concept_name_to_intervene> concept group separately and then summing the
  expectations.

  Args:
    intervened_ds: The IntervenedDataset object to use for metric computation.
    concept_name_to_intervene: The name of the concept group to intervene on in
      the next intervention step.
    curr_intervention_mask: Intervention mask being used for the current
      intervention step.
    signed: Whether to use the signed difference for entropy to measure change.

  Returns:
    The expected change in prediced label confidence.
  """
  concept_probs, dist, dists_0, dists_1 = label_metric_utilsv2(
      intervened_ds, concept_name_to_intervene, curr_intervention_mask)

  if intervened_ds.n_classes == 1:
    pred_class = t.cast(dist.logits > 0, t.float32)
  else:
    pred_class = tf.argmax(dist.logits, axis=1)
  pred_class_prob = dist.prob(pred_class)[:, None]
  pred_class_prob_0 = tf.stack([dist_0.prob(pred_class) for dist_0 in dists_0],
                               axis=1)
  pred_class_prob_1 = tf.stack([dist_1.prob(pred_class) for dist_1 in dists_1],
                               axis=1)
  changes_0 = pred_class_prob_0 - pred_class_prob
  changes_1 = pred_class_prob_1 - pred_class_prob

  if signed:
    metric = tf.reduce_sum(
        concept_probs * changes_1 + (1 - concept_probs) * changes_0, axis=1)
  else:
    metric = tf.reduce_sum(
        (concept_probs * tf.abs(changes_1) +
         (1 - concept_probs) * tf.abs(changes_0)),
        axis=1)
  return metric
//  [python]
def label_confidence_changev2(
    intervened_ds, concept_name_to_intervene,
    curr_intervention_mask, signed):
  """Metric that measures the importance of a concept using the expected change in the predicted label confidence.

  This v2 method deals with categorical concepts differently by taking
  expectation under the distribution each predicted binary concept in the
  <concept_name_to_intervene> concept group separately and then summing the
  expectations.

  Args:
    intervened_ds: The IntervenedDataset object to use for metric computation.
    concept_name_to_intervene: The name of the concept group to intervene on in
      the next intervention step.
    curr_intervention_mask: Intervention mask being used for the current
      intervention step.
    signed: Whether to use the signed difference for entropy to measure change.

  Returns:
    The expected change in prediced label confidence.
  """
  concept_probs, dist, dists_0, dists_1 = label_metric_utilsv2(
      intervened_ds, concept_name_to_intervene, curr_intervention_mask)

  if intervened_ds.n_classes == 1:
    pred_class = tf.cast(dist.logits > 0, tf.float32)
  else:
    pred_class = tf.argmax(dist.logits, axis=1)
  pred_class_prob = dist.prob(pred_class)[:, None]
  pred_class_prob_0 = tf.stack([dist_0.prob(pred_class) for dist_0 in dists_0],
                               axis=1)
  pred_class_prob_1 = tf.stack([dist_1.prob(pred_class) for dist_1 in dists_1],
                               axis=1)
  changes_0 = pred_class_prob_0 - pred_class_prob
  changes_1 = pred_class_prob_1 - pred_class_prob

  if signed:
    metric = tf.reduce_sum(
        concept_probs * changes_1 + (1 - concept_probs) * changes_0, axis=1)
  else:
    metric = tf.reduce_sum(
        (concept_probs * tf.abs(changes_1) +
         (1 - concept_probs) * tf.abs(changes_0)),
        axis=1)
  return metric
//  [python]
def corr_block(corr_pyramid_inst, coords, radius):
  r = int(radius)
  b, h1, w1, _ = tf.unstack(tf.shape(coords))
  out_pyramid = []
  for i, corr in enumerate(corr_pyramid_inst):
    start = t.float32(-r)
    stop = t.float32(r)
    num = t.int32(2 * r + 1)

    dx = tf.linspace(start, stop, num)
    dy = tf.linspace(start, stop, num)
    delta = tf.stack(tf.meshgrid(dy, dx), axis=-1)

    centroid_lvl = tf.reshape(coords, (b * h1 * w1, 1, 1, 2)) / 2**i
    delta_lvl = tf.reshape(delta, (1, 2 * r + 1, 2 * r + 1, 2))
    coords_lvl = t.float32(
        centroid_lvl) + t.float32(
            delta_lvl)

    corr = smurf_utils.resampler(corr, coords_lvl)

    channel_dim = (2 * r + 1) * (2 * r + 1)
    corr = tf.reshape(corr, (b, h1, w1, channel_dim))
    out_pyramid.append(corr)
  out = tf.concat(out_pyramid, axis=-1)
  return out
//  [python]
def corr_block(corr_pyramid_inst, coords, radius):
  r = int(radius)
  b, h1, w1, _ = tf.unstack(tf.shape(coords))
  out_pyramid = []
  for i, corr in enumerate(corr_pyramid_inst):
    start = tf.cast(-r, dtype=tf.float32)
    stop = tf.cast(r, dtype=tf.float32)
    num = tf.cast(2 * r + 1, tf.int32)

    dx = tf.linspace(start, stop, num)
    dy = tf.linspace(start, stop, num)
    delta = tf.stack(tf.meshgrid(dy, dx), axis=-1)

    centroid_lvl = tf.reshape(coords, (b * h1 * w1, 1, 1, 2)) / 2**i
    delta_lvl = tf.reshape(delta, (1, 2 * r + 1, 2 * r + 1, 2))
    coords_lvl = tf.cast(
        centroid_lvl, dtype=tf.float32) + tf.cast(
            delta_lvl, dtype=tf.float32)

    corr = smurf_utils.resampler(corr, coords_lvl)

    channel_dim = (2 * r + 1) * (2 * r + 1)
    corr = tf.reshape(corr, (b, h1, w1, channel_dim))
    out_pyramid.append(corr)
  out = tf.concat(out_pyramid, axis=-1)
  return out
//  [python]
def __init__(self, d_k, att_dropout=0.1):
        super(MixtureOfSoftMaxACF, self).__init__()
        self.temperature = tf.math.pow(tf.cast(d_k, tf.float32), 0.5)
        self.att_dropout = att_dropout
        
        self.dropout = tf.keras.layers.Dropout(att_dropout)
        self.softmax_1 = tf.keras.layers.Activation("softmax")
        self.softmax_2 = tf.keras.layers.Activation("softmax")

        self.d_k = d_k
//  [python]
def call(self, X):

        #X[:, :, 0] - categorical index of the element type
        Xid = tf.to_float(tf.one_hot(tf.to_int32(X[:, :, 0]), self.num_input_classes))

        #X[:, :, 1:] - all the other non-categorical features
        Xprop = X[:, :, 1:]
        return tf.concat([Xid, Xprop], axis=-1)
//  [python]
def call(self, X):

        #X[:, :, 0] - categorical index of the element type
        Xid = tf.cast(tf.one_hot(tf.cast(X[:, :, 0], tf.int32), self.num_input_classes), dtype=tf.float32)

        #X[:, :, 1:] - all the other non-categorical features
        Xprop = X[:, :, 1:]
        return tf.concat([Xid, Xprop], axis=-1)
//  [python]
def sparse_dense_matmult_batch(sp_a, b):

    num_batches = tf.shape(b)[0]
    def map_function(x):
        i, dense_slice = x[0], x[1]
        num_points = tf.shape(b)[1]

        sparse_slice = tf.sparse.reshape(tf.sparse.slice(
            sp_a, [i, 0, 0], [1, num_points, num_points]),
            [num_points, num_points])
        mult_slice = tf.sparse.sparse_dense_matmul(sparse_slice, dense_slice)
        return mult_slice

    elems = (tf.range(0, t.convert_to_tensor(num_batches, dtype=tf.int64), delta=1, dtype=tf.int64), b)
    ret = tf.map_fn(map_function, elems, fn_output_signature=tf.float32, back_prop=True)
    return ret
//  [python]
def sparse_dense_matmult_batch(sp_a, b):

    num_batches = tf.shape(b)[0]
    def map_function(x):
        i, dense_slice = x[0], x[1]
        num_points = tf.shape(b)[1]

        sparse_slice = tf.sparse.reshape(tf.sparse.slice(
            sp_a, [i, 0, 0], [1, num_points, num_points]),
            [num_points, num_points])
        mult_slice = tf.sparse.sparse_dense_matmul(sparse_slice, dense_slice)
        return mult_slice

    elems = (tf.range(0, tf.cast(num_batches, tf.int64), delta=1, dtype=tf.int64), b)
    ret = tf.map_fn(map_function, elems, fn_output_signature=tf.float32, back_prop=True)
    return ret
//  [python]
def train(args):
    training_samples = get_samples_from_folder(args.training_data)
    allreduce_controller = create_elastic_controller(
        batch_size=args.batch_size,
        num_epochs=args.num_epochs,
        dataset_size=len(training_samples),
        shuffle=True,
    )

    # Create the dataset with the elastic allreduce controller
    # which will fetch data shards from the master of ElasticDL job.
    dataset = create_dataset(
        allreduce_controller.data_shard_service, training_samples
    )
    dataset = dataset.batch(args.batch_size).prefetch(1)
    dataset_it = dataset.make_one_shot_iterator()
    batch_x, batch_y = dataset_it.get_next()
    batch_x = tf.cast(batch_x, tf.float32)

    batch_y = tf.reshape(batch_y, (-1,))
    predict, loss = conv_model(batch_x, batch_y, tf.estimator.ModeKeys.TRAIN)
    optimizer = tf.train.GradientDescentOptimizer(0.1)

    # Wrap the optimizer.
    optimizer = DistributedOptimizer(optimizer, fixed_global_batch_size=True)
    global_step = tf.train.get_or_create_global_step()
    train_step = optimizer.minimize(loss, global_step=global_step)

    # Use the elastic wrapper to wrap the function to train one batch
    elastic_train_one_batch = allreduce_controller.elastic_run(train_one_batch)
    hook = AdjustBackwardPassesPerStepHook(optimizer)
    allreduce_controller.set_broadcast_variables(tf.global_variables())
    with allreduce_controller.scope():
        with tf.train.MonitoredTrainingSession(hooks=[hook]) as sess:
            allreduce_controller.set_session(sess)
            try:
                while True:
                    loss_value, step, _ = elastic_train_one_batch(
                        sess, [loss, global_step, train_step]
                    )
                    logger.info(
                        "global step = {}. loss: {}".format(step, loss_value)
                    )
            except tf.errors.OutOfRangeError:
                print("end!")
//  [python]
def separate_truth(y_true):
    true_id = tf.to_int32(y_true[:, :, :1])
    true_momentum = y_true[:, :, 1:4]
    true_charge = y_true[:, :, 4:5]
    return true_id, true_charge, true_momentum
//  [python]
def separate_truth(y_true):
    true_id = tf.cast(y_true[:, :, :1], tf.int32)
    true_momentum = y_true[:, :, 1:4]
    true_charge = y_true[:, :, 4:5]
    return true_id, true_charge, true_momentum
//  [python]
def my_loss_cls(y_true, y_pred):
    pred_id_logits, pred_charge, _ = separate_prediction(y_pred)
    true_id, true_charge, _ = separate_truth(y_true)

    true_id_onehot = tf.one_hot(tf.to_int32(true_id), depth=len(class_labels))
    #predict the particle class labels
    l1 = mult_classification_loss*tf.nn.softmax_cross_entropy_with_logits(true_id_onehot, pred_id_logits)
    l3 = mult_charge_loss*mse_unreduced(true_charge, pred_charge)[:, :, 0]

    loss = l1 + l3
    return mult_total_loss*loss
//  [python]
def my_loss_cls(y_true, y_pred):
    pred_id_logits, pred_charge, _ = separate_prediction(y_pred)
    true_id, true_charge, _ = separate_truth(y_true)

    true_id_onehot = tf.one_hot(tf.cast(true_id, tf.int32), depth=len(class_labels))
    #predict the particle class labels
    l1 = mult_classification_loss*tf.nn.softmax_cross_entropy_with_logits(true_id_onehot, pred_id_logits)
    l3 = mult_charge_loss*mse_unreduced(true_charge, pred_charge)[:, :, 0]

    loss = l1 + l3
    return mult_total_loss*loss