class_docstr: Test DeformableDETRModel loss test.
class_docstr_tok: ["Test","DeformableDETRModel","loss","test","."]
code: def setUp(self):
      super().setUp()
  
      self.num_classes = 5
      self.input_shape = (3, 128, 128, 3)
      self.num_decoder_layers = 2
      config = mini_config.get_config()
  
      # Create and initialize the model.
      model_cls = DeformableDETRModel
      self.model = model_cls(
          config=config,
          dataset_meta_data={
              "num_classes": self.num_classes,
              "target_is_onehot": False,
          },
      )
  
      rng = random.PRNGKey(0)
      initial_params = self.model.flax_model.init(
          rng, jnp.zeros(self.input_shape, jnp.float32), train=False
      )
      flax_model = functools.partial(
          self.model.flax_model.apply, initial_params
      )
  
      # A fake batch with 3 examples.
      self.batch = {
          "inputs": jnp.array(
              np.random.normal(size=self.input_shape)
          ).astype(jnp.float32),
          "padding_mask": jnp.array(
              np.random.normal(size=self.input_shape[:-1])
          ).astype(jnp.float32),
          "label": {
              "labels": jnp.array(
                  np.random.randint(
                      self.num_classes,
                      size=(3, self.model.config.num_queries),
                  )
              ),
              "boxes": jnp.array(
                  np.random.uniform(
                      size=(3, self.model.config.num_queries, 4),
                      low=0.0,
                      high=1.0,
                  ),
                  dtype=jnp.float32,
              ),
          },
      }
      self.outputs = flax_model(
          self.batch["inputs"],
          padding_mask=self.batch["padding_mask"],
          train=False,
      )
  
      seq = np.arange(self.model.config.num_queries, dtype=np.int32)
      seq_rev = seq[::-1]
      seq_21 = np.concatenate(
          [
              seq[self.model.config.num_queries // 2 :],
              seq[: self.model.config.num_queries // 2],
          ]
      )
      self.indices = jnp.array(
          [(seq, seq_rev), (seq_rev, seq), (seq, seq_21)]
      )
  
file_path: scenic/scenic/projects/baselines/deformable_detr/tests/test_model.py
imports: import functools
  from absl.testing import absltest
  from absl.testing import parameterized
  from flax import jax_utils
  import jax
  from jax import random
  import jax.numpy
  import numpy
  from scenic.projects.baselines.deformable_detr.configs import mini_config
  from scenic.projects.baselines.deformable_detr.model import compute_cost
  from scenic.projects.baselines.deformable_detr.model import DeformableDETRModel
label_desc: Use image node_id 1 for calling the DeformableDETRModelLossTest obj's underlying member method code with example usage: obj.setUp() without return types