class_docstr: Unit tests for classes from module.py.
class_docstr_tok: ["Unit","tests","for","classes","from","module.py","."]
code: def test_force_args(self):
      # def __call__(self, x, some_arg1, some_arg2):
      #     # without forces args equivalent to (x + some_arg1 * 5 + some_arg2 * 2)
      x = jn.array([1.0, 2.0])
      model = ModelWithArg()
      np.testing.assert_almost_equal(
          model(x, some_arg1=0.0, some_arg2=0.0), [1.0, 2.0]
      )
      np.testing.assert_almost_equal(
          model(x, some_arg1=1.0, some_arg2=0.0), [6.0, 7.0]
      )
      np.testing.assert_almost_equal(
          model(x, some_arg1=1.0, some_arg2=1.0), [8.0, 9.0]
      )
      # Ensure that can not use invalid argument in original module
      with self.assertRaises(TypeError):
          model(x, some_arg1=0.0, some_arg2=0.0, wrong_arg_name=0.0)
      with self.assertRaises(TypeError):
          model(x, wrong_arg_name1=0.0, wrong_arg_name2=0.0)
      with self.assertRaises(TypeError):
          model.block1.op1(x, wrong_arg_name=0.0)
      # Set forced args.
      original_signature = inspect.signature(model.block1.op1)
      model.block1.op1 = objax.ForceArgs(
          model.block1.op1, some_arg=-1.0
      )
      # At this point following arguments are forced:
      #   model.block1.op1(..., some_arg=-1.0)
      self.assertEqual(
          original_signature, inspect.signature(model.block1.op1)
      )
      np.testing.assert_almost_equal(
          model(x, some_arg1=0.0, some_arg2=0.0), [0.0, 1.0]
      )
      np.testing.assert_almost_equal(
          model(x, some_arg1=1.0, some_arg2=0.0), [5.0, 6.0]
      )
      np.testing.assert_almost_equal(
          model(x, some_arg1=1.0, some_arg2=1.0), [6.0, 7.0]
      )
      # ForceArgs does not allow to pass invalid args
      with self.assertRaises(TypeError):
          model.block1.op1(x, wrong_arg_name=1.0)
      # Set forced args in a list
      model.block1.seq[0] = objax.ForceArgs(
          model.block1.seq[0], some_arg=-1.0
      )
      # At this point following arguments are forced:
      #   model.block1.op1(..., some_arg=-1.0)
      #   model.block1.seq[0](..., some_arg=-1.0)
      np.testing.assert_almost_equal(
          model(x, some_arg1=0.0, some_arg2=0.0), [-1.0, 0.0]
      )
      np.testing.assert_almost_equal(
          model(x, some_arg1=1.0, some_arg2=0.0), [3.0, 4.0]
      )
      np.testing.assert_almost_equal(
          model(x, some_arg1=1.0, some_arg2=1.0), [4.0, 5.0]
      )
      # Set invalid arg in forced args
      model.block2.op1 = objax.ForceArgs(
          model.block2.op1, wrong_arg_name=1.0
      )
      with self.assertRaises(TypeError):
          model(x, some_arg=0.0)
      # Remove force args with invalid name
      objax.ForceArgs.undo(model, wrong_arg_name=objax.ForceArgs.ANY)
      np.testing.assert_almost_equal(
          model(x, some_arg1=0.0, some_arg2=0.0), [-1.0, 0.0]
      )
      # Set few other forces args with nested ForceArgs
      model.block1 = objax.ForceArgs(model.block1, some_arg1=10.0)
      model.block1 = objax.ForceArgs(model.block1, some_arg2=20.0)
      # At this point following arguments are forced:
      #   model.block1(..., some_arg1=10.0, some_arg2=20.0)
      #   model.block1.op1(..., some_arg=-1.0)
      #   model.block1.seq[0](..., some_arg=-1.0)
      np.testing.assert_almost_equal(
          model(x, some_arg1=0.0, some_arg2=0.0), [9.0, 10.0]
      )
      np.testing.assert_almost_equal(
          model(x, some_arg1=1.0, some_arg2=0.0), [12.0, 13.0]
      )
      np.testing.assert_almost_equal(
          model(x, some_arg1=1.0, some_arg2=1.0), [13.0, 14.0]
      )
      # Resetting some of the forced args
      objax.ForceArgs.undo(
          model, some_arg1=30
      )  # noop because some_arg1=30 is not used
      np.testing.assert_almost_equal(
          model(x, some_arg1=0.0, some_arg2=0.0), [9.0, 10.0]
      )
      np.testing.assert_almost_equal(
          model(x, some_arg1=1.0, some_arg2=0.0), [12.0, 13.0]
      )
      np.testing.assert_almost_equal(
          model(x, some_arg1=1.0, some_arg2=1.0), [13.0, 14.0]
      )
      objax.ForceArgs.undo(model, some_arg1=10)  # undo some_arg1=10
      np.testing.assert_almost_equal(
          model(x, some_arg1=0.0, some_arg2=0.0), [-1.0, 0.0]
      )
      np.testing.assert_almost_equal(
          model(x, some_arg1=1.0, some_arg2=0.0), [3.0, 4.0]
      )
      np.testing.assert_almost_equal(
          model(x, some_arg1=1.0, some_arg2=1.0), [4.0, 5.0]
      )
      objax.ForceArgs.undo(model)  # undo all forced args
      np.testing.assert_almost_equal(
          model(x, some_arg1=0.0, some_arg2=0.0), [1.0, 2.0]
      )
      np.testing.assert_almost_equal(
          model(x, some_arg1=1.0, some_arg2=0.0), [6.0, 7.0]
      )
      np.testing.assert_almost_equal(
          model(x, some_arg1=1.0, some_arg2=1.0), [8.0, 9.0]
      )
      # Try forced args on the root
      model = objax.ForceArgs(model, some_arg1=-1.0)
      np.testing.assert_almost_equal(
          model(x, some_arg1=0.0, some_arg2=0.0), [-4.0, -3.0]
      )
      np.testing.assert_almost_equal(
          model(x, some_arg1=1.0, some_arg2=0.0), [-4.0, -3.0]
      )
      np.testing.assert_almost_equal(
          model(x, some_arg1=1.0, some_arg2=1.0), [-2.0, -1.0]
      )
      # Reset forced args on the model root
      objax.ForceArgs.undo(model)
      np.testing.assert_almost_equal(
          model(x, some_arg1=0.0, some_arg2=0.0), [1.0, 2.0]
      )
      np.testing.assert_almost_equal(
          model(x, some_arg1=1.0, some_arg2=0.0), [6.0, 7.0]
      )
      np.testing.assert_almost_equal(
          model(x, some_arg1=1.0, some_arg2=1.0), [8.0, 9.0]
      )
  
imports: import functools
  import inspect
  import unittest
  from typing import Tuple
  import numpy
  from jax import numpy
  import objax
  from objax.typing import JaxArray
label_desc: Use image node_id 6 for calling the TestModule obj's underlying member method code with example usage: obj.test_force_args() without return types