def testCreateCanvasAndTargets(self):
    with self.session():
      tf.random.set_seed(_TF_RANDOM_SEED)
      batch = py_utils.NestedMap(
          src=py_utils.NestedMap(
              ids=tf.convert_to_tensor(
                  np.asarray([
                      [10, 11, 12, 14, 2, 0],
                      [20, 21, 22, 24, 25, 2],
                  ], np.int32)),
              paddings=tf.convert_to_tensor(
                  np.asarray([[0, 0, 0, 0, 0, 1], [0, 0, 0, 0, 0, 0]],
                             np.float32))),
          tgt=py_utils.NestedMap(
              ids=tf.convert_to_tensor(
                  np.asarray([[100, 101, 102, 104, 2, 0],
                              [200, 201, 202, 204, 205, 2]], np.int32)),
              paddings=tf.convert_to_tensor(
                  np.asarray([[0, 0, 0, 0, 0, 1], [0, 0, 0, 0, 0, 0]],
                             np.float32))))

      p = self._testParams()
      mdl = p.Instantiate()

      descriptor = mdl._CreateCanvasAndTargets(batch)

      canvas, canvas_paddings, target_indices, target_weights = self.evaluate([
          descriptor.canvas, descriptor.canvas_paddings,
          descriptor.target_indices, descriptor.target_weights
      ])

      canvas_gold = np.asarray([
          [32014, 32002, 104, 2, 0, 0, 0, 0],
          [32020, 32021, 32022, 32002, 200, 201, 202, 2],
      ], np.int32)
      canvas_paddings_gold = np.asarray(
          [[0., 0., 0., 0., 1., 1., 1., 1.], [0., 0., 0., 0., 0., 0., 0., 0.]],
          np.float32)
      target_indices_gold = np.asarray(
          [[0, 0, 10], [0, 0, 11], [0, 0, 12], [0, 0, 2], [0, 1, 2], [1, 0, 2],
           [1, 1, 2], [1, 2, 2], [1, 3, 24], [1, 3, 25], [1, 3, 2], [0, 2, 100],
           [0, 2, 101], [0, 2, 102], [0, 2, 2], [0, 3, 2], [1, 4, 2], [1, 5, 2],
           [1, 6, 2], [1, 7, 204], [1, 7, 205], [1, 7, 2]], np.int32)
      target_weights_gold = np.asarray([1, 1, 1, 0, 1] + [1, 1, 1, 1, 1, 0] +
                                       [1, 1, 1, 0, 1] + [1, 1, 1, 1, 1, 0],
                                       np.float32)
      target_weights_gold = np.reshape(target_weights_gold,
                                       [target_weights_gold.shape[0], 1])

      self.assertAllEqual(canvas, canvas_gold)
      self.assertAllEqual(canvas_paddings, canvas_paddings_gold)
      self.assertAllEqual(target_indices, target_indices_gold)
      self.assertAllEqual(target_weights, target_weights_gold)