class_docstr: Test train utilities.
class_docstr_tok: ["Test","train","utilities","."]
docstr: Setup sample output predictions and target labels and bounding boxes.
docstr_tok: ["Setup","sample","output","predictions","and","target","labels","and","bounding","boxes","."]
code: def setUp(self):
      """Setup sample output predictions and target labels and bounding boxes."""
      super().setUp()
  
      self.batchsize = 4
      self.num_classes = 81
      self.max_num_boxes = 63
      self.num_preds = 100
  
      key = jax.random.PRNGKey(0)
  
      # Create fake predictions and targets
      key, subkey = jax.random.split(key)
      # set probabilities for class 0 higher than others
      p_logits = jnp.ones(self.num_classes).at[0].set(5.0)
      p = jax.nn.softmax(p_logits)
      tgt_labels = jax.random.choice(
          subkey,
          self.num_classes,
          shape=(self.batchsize, self.max_num_boxes),
          replace=True,
          p=p,
      )
      # Ensure last target is dummy empty target.
      tgt_labels = tgt_labels.at[:, -1].set(0)
      onehot_tgt_labels = jax.nn.one_hot(tgt_labels, self.num_classes)
  
      key, subkey = jax.random.split(key)
      pred_logits = jax.random.normal(
          subkey,
          shape=(self.batchsize, self.num_preds, self.num_classes),
      )
      pred_probs = jax.nn.softmax(pred_logits, axis=-1)
  
      key, subkey = jax.random.split(key)
      pred_bbox = sample_cxcywh_bbox(
          subkey, batch_shape=(self.batchsize, self.num_preds)
      )
  
      key, subkey = jax.random.split(key)
      tgt_bbox = sample_cxcywh_bbox(
          subkey, batch_shape=(self.batchsize, self.max_num_boxes)
      )
  
      self.outputs = {"pred_probs": pred_probs, "pred_boxes": pred_bbox}
      self.targets = {"labels": tgt_labels, "boxes": tgt_bbox}
      self.onehot_targets = {
          "labels": onehot_tgt_labels,
          "boxes": tgt_bbox,
      }
  
file_path: scenic/scenic/projects/baselines/detr/tests/test_train_utils.py
imports: import collections
  import os
  from absl.testing import absltest
  from absl.testing import parameterized
  import jax
  import jax.numpy
  import numpy
  from scenic.projects.baselines.detr import train_utils
label_desc: Use image node_id 1 for calling the TrainUtilsTest obj's underlying member method code with example usage: obj.setUp() without return types