code: def test_on_linear(self):
      k = objax.nn.Linear(3, 3)
      kj = objax.Jit(k)
      x = objax.random.normal((64, 3))
      y1 = kj(x)
      k.w.assign(k.w.value + 1)
      y2 = kj(x)
      k.w.assign(k.w.value - 1)
      y3 = kj(x)
      self.assertAlmostEqual(((y1 - y3) ** 2).sum(), 0)
      self.assertNotEqual(((y1 - y2) ** 2).sum(), 0)
  
file_path: objax/tests/jit.py
imports: import unittest
  import jax.numpy
  from jax.core import ConcretizationTypeError
  import objax
  from objax.typing import JaxArray
label_desc: Use image node_id 1 for calling the TestJit obj's underlying member method code with example usage: obj.test_on_linear() without return types