Add cache tests

examples/ffi/tests/attrs_test.py

@@ -28,6 +28,16 @@ def test_array_attr(self):
  self.assertEqual(attrs.array_attr(5), jnp.arange(5).sum())
  self.assertEqual(attrs.array_attr(3), jnp.arange(3).sum())
 
+ def test_array_attr_jit_cache(self):
+ jit_array_attr = jax.jit(attrs.array_attr, static_argnums=(0,))
+ with jtu.count_jit_and_pmap_lowerings() as count:
+ jit_array_attr(5)
+ self.assertEqual(count[0], 1) # compiles once the first time
+ with jtu.count_jit_and_pmap_lowerings() as count:
+ jit_array_attr(5)
+ self.assertEqual(count[0], 0) # cache hit
+ self.assertNotIn("_HashableByObjectId", jit_array_attr.lower(5).as_text())
+
  def test_dictionary_attr(self):
  secret, count = attrs.dictionary_attr(secret=5)
  self.assertEqual(secret, 5)