Add jit unit test to autoscale decorator. (#12)

Making sure we cover eager + jit modes in our testing.

tests/core/test_interpreter.py

@@ -10,11 +10,13 @@
 
 
 class AutoScaleInterpreterTests(chex.TestCase):
- def test__identity(self):
+ @chex.variants(with_jit=True, without_jit=True)
+ def test__scaled_identity_function(self):
  def func(x):
  return x
 
- asfunc = autoscale(func)
+ # Autoscale + (optional) jitting.
+ asfunc = self.variant(autoscale(func))
 
  scaled_inputs = scaled_array([1.0, 2.0], 1, dtype=np.float32)
  scaled_outputs = asfunc(scaled_inputs)
@@ -31,11 +33,13 @@ def func(x):
  assert jaxpr.outvars[0].aval.shape == expected.shape
  assert jaxpr.outvars[1].aval.shape == ()
 
- def test__mul(self):
+ @chex.variants(with_jit=True, without_jit=True)
+ def test__scaled_mul_function(self):
  def func(x, y):
  return x * y
 
- asfunc = autoscale(func)
+ # Autoscale + (optional) jitting.
+ asfunc = self.variant(autoscale(func))
 
  x = scaled_array([-2.0, 2.0], 0.5, dtype=np.float32)
  y = scaled_array([1.5, 1.5], 2, dtype=np.float32)