code: def setUp(self):
      super().setUp()
      self.class_logits = np.array(
          [
              [
                  [[1.2], [0.4], [-0.9]],
                  [[-0.4], [0.8], [0.1]],
                  [[-1.0], [-1.0], [-1.0]],
                  [[-1.0], [-1.0], [-1.0]],
              ],
              [
                  [[-1.0], [-1.0], [-1.0]],
                  [[-1.0], [-1.0], [-1.0]],
                  [[1.2], [0.9], [0.4]],
                  [[0.1], [0.1], [0.1]],
              ],
          ]
      )  # shape is (2, 2*2, 3, 1).
      pred_displacements = np.array(
          [
              [
                  [[0.4, 0.4], [0.0, 0.1], [0.2, 0.0]],
                  [[0.1, 0.2], [0.3, 0.4], [0.5, 0.6]],
                  [[0.0, 0.0], [0.0, 0.0], [0.0, 0.0]],
                  [[0.0, 0.0], [0.0, 0.0], [0.0, 0.0]],
              ],
              [
                  [[0.0, 0.0], [0.0, 0.0], [0.0, 0.0]],
                  [[0.0, 0.0], [0.0, 0.0], [0.0, 0.0]],
                  [[1.2, 0.9], [0.4, 0.4], [0.4, 0.2]],
                  [[0.4, 0.8], [0.1, 0.3], [0.3, 0.3]],
              ],
          ]
      )  # shape is (2, 2*2, 3, 2).
      self.logits = np.concatenate(
          [self.class_logits, pred_displacements], axis=-1
      )
      self.batch = {
          "batch_mask": np.ones((2,), dtype=np.int32),
          "inputs": {
              "input_mask": np.array(
                  [[1, 1, 1], [1, 1, 0]], dtype=np.int32
              ),
              "caption_mask": np.array(
                  [[1, 1], [1, 0]], dtype=np.int32
              ),
          },
          "label": np.array(
              [
                  [[[1], [1], [0]], [[0], [1], [1]]],
                  [[[1], [1], [0]], [[0], [0], [0]]],
              ],
              dtype=np.int32,
          ),  # shape is (2, 2, 3, 1).
          "displacements": np.array(
              [
                  [
                      [[0.5, 0.4], [0.3, 0.2], [0, 0]],
                      [[0, 0], [0.4, 0.3], [0.5, 0.6]],
                  ],
                  [
                      [[0.2, 0.1], [0.3, 0.3], [0, 0]],
                      [[0, 0], [0, 0], [0, 0]],
                  ],
              ],
              dtype=np.float32,
          ),  # shape is (2, 2, 3, 2).
      }
  
file_path: scenic/scenic/projects/unloc/moment_retrieval_base_model_test.py
imports: from absl.testing import absltest
  from absl.testing import parameterized
  from flax import jax_utils
  import jax
  import ml_collections
  import numpy
  from scenic.projects.unloc import moment_retrieval_base_model
label_desc: Use image node_id 1 for calling the MomentRetrievalBaseModelTest obj's underlying member method code with example usage: obj.setUp() without return types