class_docstr: Tests the DeformableDETR trainer on single device setup.
class_docstr_tok: ["Tests","the","DeformableDETR","trainer","on","single","device","setup","."]
docstr: Test DeformableDETR eval_step directly.
docstr_tok: ["Test","DeformableDETR","eval_step","directly","."]
code: def test_deformable_detr_eval_step(self):
      """Test DeformableDETR eval_step directly."""
      rng = jax.random.PRNGKey(0)
      np.random.seed(0)
  
      (
          model,
          _,
          train_state,
          _,
          _,
          _,
      ) = trainer.get_model_and_tx_and_train_state(
          rng=rng,
          dataset=self.dataset,
          config=self.config,
          model_cls=self.model_cls,
          workdir="",
          input_spec=[
              (
                  self.dataset.meta_data["input_shape"],
                  self.dataset.meta_data.get(
                      "input_dtype", jnp.float32
                  ),
              )
          ],
      )
  
      eval_step = ddetr_eval.get_eval_step(
          flax_model=model.flax_model,
          loss_and_metrics_fn=model.loss_and_metrics_function,
          logits_to_probs_fn=model.logits_to_probs,
          debug=False,
      )
      eval_step_pmapped = jax.pmap(
          eval_step, axis_name="batch", donate_argnums=(1,)
      )
  
      with tfds.testing.mock_data(
          num_examples=50,
          as_dataset_fn=test_util.generate_fake_dataset(
              num_examples=50
          ),
      ):
          eval_batch = next(self.dataset.valid_iter)
  
      init_params, _ = flatten_util.ravel_pytree(train_state.params)
      _, predictions, metrics = eval_step_pmapped(
          train_state, eval_batch
      )
      after_params, _ = flatten_util.ravel_pytree(train_state.params)
      self.assertTrue(jnp.array_equal(init_params, after_params))
  
      exp_pred_keys = {"pred_logits", "pred_boxes", "pred_probs"}
      self.assertSameElements(predictions.keys(), exp_pred_keys)
      exp_metric_keys = [
          "loss_bbox",
          "loss_class",
          "loss_giou",
          "total_loss",
      ]
      for i in range(self.config.num_decoder_layers - 1):
          aux = [
              f"loss_class_aux{i}",
              f"loss_bbox_aux{i}",
              f"loss_giou_aux{i}",
          ]
          exp_metric_keys += aux
      self.assertSameElements(exp_metric_keys, metrics.keys())
  
imports: import shutil
  import tempfile
  from absl.testing import absltest
  from absl.testing import parameterized
  from jax import flatten_util
  import jax.numpy
  import jax.random
  import numpy
  from scenic.dataset_lib import datasets
  from scenic.projects.baselines.deformable_detr import evaluate
  from scenic.projects.baselines.deformable_detr import input_pipeline_detection
  from scenic.projects.baselines.deformable_detr import trainer
  from scenic.projects.baselines.deformable_detr.configs import mini_config
  from scenic.projects.baselines.deformable_detr.model import DeformableDETRModel
  from scenic.projects.baselines.detr.tests import test_util
  import tensorflow
  import tensorflow_datasets
label_desc: Use image node_id 3 for calling the DETRTrainerTest obj's underlying member method code with example usage: obj.test_deformable_detr_eval_step() without return types