class_docstr: Test DeformableDETRModel cost test.
class_docstr_tok: ["Test","DeformableDETRModel","cost","test","."]
docstr: Test compute_cost.
docstr_tok: ["Test","compute_cost","."]
code: def test_compute_cost(self):
      """Test compute_cost."""
  
      bs = 3
      num_classes = 7
      num_preds = 5
      # Includes padding.
      max_targets = 10
  
      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(num_classes).at[0].set(5.0)
      tgt_labels = jax.random.choice(
          subkey,
          np.arange(1, num_classes + 1),
          shape=(bs, max_targets),
          replace=True,
          p=jax.nn.softmax(p_logits),
      )
  
      # Set padding for targets by index of batch
      for i in range(1, bs + 1):
          tgt_labels = tgt_labels.at[i, -i:].set(0)
  
      key, subkey = jax.random.split(key)
      pred_logits = jax.random.normal(
          subkey, shape=(bs, num_preds, num_classes)
      )
      pred_probs = jax.nn.sigmoid(pred_logits)
  
      key, subkey = jax.random.split(key)
      pred_bbox = sample_cxcywh_bbox(
          subkey, batch_shape=(bs, num_preds)
      )
  
      key, subkey = jax.random.split(key)
      tgt_bbox = sample_cxcywh_bbox(
          subkey, batch_shape=(bs, max_targets)
      )
  
      cost, n_cols = compute_cost(
          tgt_bbox=tgt_bbox,
          tgt_labels=tgt_labels,
          out_bbox=pred_bbox,
          out_prob=pred_probs,
          bbox_loss_coef=1.0,
          giou_loss_coef=1.0,
          class_loss_coef=1.0,
          target_is_onehot=False,
      )
      self.assertSequenceEqual(cost.shape, (bs, num_preds, max_targets))
      exp_n_cols = range(max_targets, max_targets - bs, -1)
      self.assertSequenceEqual(n_cols.tolist(), exp_n_cols)
  
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 DeformableDETRModelCostTest obj's underlying member method code with example usage: obj.test_compute_cost() without return types