class_docstr: Unit tests for classes from module.py.
class_docstr_tok: ["Unit","tests","for","classes","from","module.py","."]
docstr: Unit test for objax.Function.with_args
docstr_tok: ["Unit","test","for","objax.Function.with_args"]
code: def test_module_function_decorator(self):
      """Unit test for objax.Function.with_args"""
      m = ComplexModule(2)
  
      @objax.Function.with_vars(m.vars())
      def my_func(
          x: JaxArray, y: JaxArray, *, named: int = 1
      ) -> Tuple[JaxArray, JaxArray]:
          return (
              m(x, training=True) - named,
              m(y, training=False) + named,
          )
  
      self.assertEqual(my_func.vars(), m.vars("{my_func}"))
      self.assertEqual(
          inspect.signature(my_func),
          inspect.Signature(
              [
                  inspect.Parameter(
                      "x",
                      inspect.Parameter.POSITIONAL_OR_KEYWORD,
                      annotation=JaxArray,
                  ),
                  inspect.Parameter(
                      "y",
                      inspect.Parameter.POSITIONAL_OR_KEYWORD,
                      annotation=JaxArray,
                  ),
                  inspect.Parameter(
                      "named",
                      inspect.Parameter.KEYWORD_ONLY,
                      annotation=int,
                      default=1,
                  ),
              ],
              return_annotation=Tuple[JaxArray, JaxArray],
          ),
      )
  
      x = jn.arange(6).reshape((3, 2))
      u, v = my_func(x, x + 1)
      self.assertEqual(
          u.tolist(), [[-1.0, 0.0], [1.0, 2.0], [3.0, 4.0]]
      )
      self.assertEqual(
          v.tolist(), [[0.0, -1.0], [-2.0, -3.0], [-4.0, -5.0]]
      )
      u, v = my_func(x, x + 1, named=2)
      self.assertEqual(
          u.tolist(), [[-2.0, -1.0], [0.0, 1.0], [2.0, 3.0]]
      )
      self.assertEqual(
          v.tolist(), [[1.0, 0.0], [-1.0, -2.0], [-3.0, -4.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 4 for calling the TestModule obj's underlying member method code with example usage: obj.test_module_function_decorator() and returns: unknown,  unknown