code: def test_constant_optimization(self):
      m = objax.nn.Linear(3, 4)
      jit_constant = objax.Jit(m, objax.VarCollection())
  
      x = objax.random.normal((10, 3))
      self.assertEqual(((m(x) - jit_constant(x)) ** 2).sum(), 0)
  
      # Modify m (which was supposed to be constant!)
      m.b.assign(m.b.value + 1)
      self.assertEqual(((m(x) - jit_constant(x)) ** 2).sum(), 40)
  
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 6 for calling the TestJit obj's underlying member method code with example usage: obj.test_constant_optimization() without return types