class_docstr: Unit tests for variable tracing using.
code: def test_typical_training_loop(self):
      # Define model and optimizer
      model = DNNet((32, 10), objax.functional.leaky_relu)
      opt = objax.optimizer.Momentum(model.vars(), nesterov=True)
  
      # Predict op
      predict_op = lambda x: objax.functional.softmax(
          model(x, training=False)
      )
  
      self.assertDictEqual(
          objax.util.find_used_variables(predict_op),
          model.vars(scope="model."),
      )
  
      # Loss function
      def loss(x, label):
          logit = model(x, training=True)
          xe_loss = objax.functional.loss.cross_entropy_logits_sparse(
              logit, label
          ).mean()
          return xe_loss
  
      self.assertDictEqual(
          objax.util.find_used_variables(loss),
          model.vars(scope="model."),
      )
  
      # Gradients and loss function
      loss_gv = objax.GradValues(
          loss, objax.util.find_used_variables(loss)
      )
  
      def train_op(x, y, learning_rate):
          grads, loss = loss_gv(x, y)
          opt(learning_rate, grads)
          return loss
  
      self.assertDictEqual(
          objax.util.find_used_variables(train_op),
          {
              **model.vars(scope="loss_gv.model."),
              **opt.vars(scope="opt."),
          },
      )
  
file_path: objax/tests/tracing.py
imports: import unittest
  import numpy
  import jax.numpy
  import objax
  from objax.zoo.dnnet import DNNet
label_desc: Use image node_id 8 for calling the TestTracing obj's underlying member method code with example usage: obj.test_typical_training_loop() and returns: xe_loss, loss