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.
docstr_tok: ["Unit","test","for","objax.Function","."]
code: def test_module_function(self):
      """Unit test for objax.Function."""
      m = ComplexModule(2)
  
      def my_func(
          x: JaxArray, y: JaxArray, *, named: int = 1
      ) -> Tuple[JaxArray, JaxArray]:
          return (
              m(x, training=True) - named,
              m(y, training=False) + named,
          )
  
      module_func = objax.Function(my_func, m.vars())
      self.assertEqual(module_func.vars(), m.vars("{my_func}"))
      self.assertEqual(
          inspect.signature(module_func), inspect.signature(my_func)
      )
      x = jn.arange(6).reshape((3, 2))
      self.assertEqual(
          [v.tolist() for v in module_func(x, x + 1)],
          [v.tolist() for v in my_func(x, x + 1)],
      )
      self.assertEqual(
          [v.tolist() for v in module_func(x, x + 1, named=2)],
          [v.tolist() for v in my_func(x, x + 1, named=2)],
      )
      # partial doesn't retain names.
      module_func = objax.Function(
          functools.partial(my_func, named=2), m.vars()
      )
      self.assertEqual(module_func.vars(), m.vars())
  
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 3 for calling the TestModule obj's underlying member method code with example usage: obj.test_module_function() and returns: unknown,  unknown