[export] Add support for serializing functions with PRNG keys as inputs/outputs

This introduces version 4 of serialization, fully backwards compatible
with versions 2 and 3.

Fixes: jax-ml#24143

Author: gnecula

jax/_src/export/serialization.fbs

@@ -45,7 +45,7 @@ enum AbstractValueKind: byte {
 }
 
 enum DType: byte {
-  // Last used id: 22
+  // Last used id: 29
   bool = 0,
   i8 = 1,
   i16 = 2,
@@ -76,6 +76,10 @@ enum DType: byte {
   f8_e5m2fnuz = 21,
   f8_e8m0fnu = 25,
   f4_e2m1fn = 26,
+
+  key_fry = 27,
+  key_rbg = 28,
+  key_unsafe_rbg = 29,
 }
 
 table AbstractValue {

jax/_src/export/serialization.py

@@ -31,6 +31,7 @@
 from jax._src import core
 from jax._src import dtypes
 from jax._src import effects
+from jax._src import prng
 from jax._src import tree_util
 from jax._src.export import serialization_generated as ser_flatbuf
 from jax._src.export import _export
@@ -48,6 +49,8 @@
 # Version 2, Dec 16th, 2023, adds the f0 dtype.
 # Version 3, October 16th, 2024, adds serialization for namedtuple and custom types
 #   This version is backwards compatible with Version 2.
+# Version 4, April 7th, 2025, adds serialization for PRNGs key types.
+#   This version is backwards compatible with Version 2 and 3.
 _SERIALIZATION_VERSION = 2
 
 def serialize(exp: _export.Exported, vjp_order: int = 0) -> bytearray:
@@ -361,6 +364,10 @@ def _deserialize_pytreedef_to_pytree(p: ser_flatbuf.PyTreeDef):
     dtypes._float8_e4m3_dtype: ser_flatbuf.DType.f8_e4m3,
     dtypes._float8_e8m0fnu_dtype: ser_flatbuf.DType.f8_e8m0fnu,
     dtypes._float4_e2m1fn_dtype: ser_flatbuf.DType.f4_e2m1fn,
+
+    prng.KeyTy(prng.prngs["threefry2x32"]): ser_flatbuf.DType.key_fry,
+    prng.KeyTy(prng.prngs["rbg"]): ser_flatbuf.DType.key_rbg,
+    prng.KeyTy(prng.prngs["unsafe_rbg"]): ser_flatbuf.DType.key_unsafe_rbg,
 }
 
 _dtype_kind_to_dtype = {

jax/_src/export/serialization_generated.py

@@ -53,16 +53,19 @@ class DType(object):
     bf16 = 14
     i4 = 15
     ui4 = 16
-    f8_e3m4 = 24
-    f8_e4m3 = 23
     f8_e4m3b11fnuz = 17
     f8_e4m3fn = 18
     f8_e4m3fnuz = 19
     f8_e5m2 = 20
     f8_e5m2fnuz = 21
     f0 = 22
+    f8_e4m3 = 23
+    f8_e3m4 = 24
     f8_e8m0fnu = 25
     f4_e2m1fn = 26
+    key_fry = 27
+    key_rbg = 28
+    key_unsafe_rbg = 29
 
 
 class ShardingKind(object):

jax/_src/prng.py

@@ -113,7 +113,7 @@ def pprint(self):
             ]))))
 
 
-prngs = {}
+prngs: dict[str, PRNGImpl] = {}
 
 def register_prng(impl: PRNGImpl):
   if impl.name in prngs:

tests/export_test.py

@@ -421,6 +421,18 @@ def f(x1, x2):
     self.assertEqual(tree_util.tree_structure(res2),
                      tree_util.tree_structure(res))
 
+  @jtu.parameterized_filterable(
+    kwargs=[dict(impl=p)
+            for p in ("rbg", "unsafe_rbg", "threefry2x32")])
+  def test_prng_keys(self, *, impl):
+
+    key = jax.random.key(42, impl=impl)
+    @jax.jit
+    def f(key):
+      return key
+    exp_f = get_exported(jax.jit(f))(key)
+    self.assertEqual(f(key), exp_f.call(key))
+
   def test_error_wrong_intree(self):
     def f(a_b_pair, *, c):
       return jnp.sin(a_b_pair[0]) + jnp.cos(a_b_pair[1]) + c