Fix test_assert_trees_all_equal_prng_keys

Converting from JAX PRNG key to numpy array is now an error. This happens when
construcing an error string in the tests.

PiperOrigin-RevId: 696046978

chex/_src/asserts.py

@@ -1578,8 +1578,17 @@ def err_msg_fn(arr_1, arr_2) -> str:
     try:
       assert_fn(arr_1, arr_2)
     except AssertionError as e:
-      return (f"{str(e)} \nOriginal dtypes: "
-              f"{np.asarray(arr_1).dtype}, {np.asarray(arr_2).dtype}")
+      dtype_1 = (
+          arr_1.dtype
+          if isinstance(arr_1, jax.Array)
+          else np.asarray(arr_1).dtype
+      )
+      dtype_2 = (
+          arr_2.dtype
+          if isinstance(arr_2, jax.Array)
+          else np.asarray(arr_1).dtype
+      )
+      return f"{str(e)} \nOriginal dtypes: {dtype_1}, {dtype_2}"
     return ""
 
   assert_trees_all_equal_comparator(cmp_fn, err_msg_fn, *trees)