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)