class_docstr: Tests the DeformableDETR trainer on single device setup.
class_docstr_tok: ["Tests","the","DeformableDETR","trainer","on","single","device","setup","."]
docstr: Test DeformableDETR train_step directly.
docstr_tok: ["Test","DeformableDETR","train_step","directly","."]
code: def test_deformable_detr_train_step(self, update_batch_stats):
      """Test DeformableDETR train_step directly."""
      rng = jax.random.PRNGKey(0)
      np.random.seed(0)
  
      (
          model,
          tx,
          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
                  ),
              )
          ],
      )
  
      train_step = trainer.get_train_step(
          apply_fn=model.flax_model.apply,
          loss_and_metrics_fn=model.loss_and_metrics_function,
          tx=tx,
          update_batch_stats=update_batch_stats,
          debug=False,
      )
  
      train_step_pmapped = jax.pmap(
          train_step,
          axis_name="batch",
          donate_argnums=(0,),
      )
  
      with tfds.testing.mock_data(
          num_examples=50,
          as_dataset_fn=test_util.generate_fake_dataset(
              num_examples=50
          ),
      ):
          train_batch = next(self.dataset.train_iter)
  
      init_params, _ = flatten_util.ravel_pytree(train_state.params)
      init_model_state, _ = flatten_util.ravel_pytree(
          train_state.model_state
      )
  
      train_state, _, _ = train_step_pmapped(train_state, train_batch)
  
      self.assertFalse(jnp.array_equal(init_params, train_state.params))
      if update_batch_stats:
          # make sure that state were updated in the train step
          self.assertFalse(
              jnp.array_equal(
                  init_model_state,
                  flatten_util.ravel_pytree(train_state.model_state)[0],
              )
          )
  
      else:
          # make sure the state were stayed frozen in the train step
          self.assertTrue(
              jnp.array_equal(
                  init_model_state,
                  flatten_util.ravel_pytree(train_state.model_state)[0],
              )
          )
  
imports: import shutil
  import tempfile
  from absl.testing import absltest
  from absl.testing import parameterized
  from clu import metric_writers
  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 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 4 for calling the DETRTrainerTest obj's underlying member method code with example usage: obj.test_deformable_detr_train_step(update_batch_stats) without return types