Repository: google/jax

Issue:
Cannot take gradient of VJP function involving new-style RNG key
### Description
```python
from jax import custom_vjp, grad
from jax.random import key

@custom_vjp
def find_fixed_point(theta, state):
    return state

def _ffp_fwd(theta, state):
    return state, None

def _ffp_bwd(residuals, state_bar):
    assert False

find_fixed_point.defvjp(_ffp_fwd, _ffp_bwd)

def fixed_point_using_while_of_theta(theta) -> float:
    state = (8.0, key(123))
    x, _ = find_fixed_point(theta, state)
    return x

grad(fixed_point_using_while_of_theta)(3.0)
```
gives
```
ValueError: Cannot convert_element_type to dtype=key<fry>
```
@froystig 

### What jax/jaxlib version are you using?

0.4.20

### Which accelerator(s) are you using?

_No response_

### Additional system info

_No response_

### NVIDIA GPU info

_No response_


Discussion:
Thanks for the report – if you could figure out where the `convert_element_type` primitive is being called, that would help. You might start by setting `JAX_TRACEBACK_FILTERING='off'` to see the full traceback.

I tried running locally to set a breakpoint, but I'm not able to import `tjax` because of Python version incompatibilities.
@jakevdp Yeah, I tried turning off the filtering, but I'm not so good at understanding the Jax internals :smile: .  I've edited the issue and removed the TJax dependency if you have time to look at it again.
Okay, I think it's close to minimal now.
It looks like RNG keys can no longer be primals to custom VJPs?
Thanks for the repro! Definitely an oversight. I think #18444 should fix it (I hit the `assert False` in your minimal example)
I just ran my full TJax tests and they all pass now.  Thanks again for the quick fix!

Resolving patch:
diff --git a/jax/_src/interpreters/ad.py b/jax/_src/interpreters/ad.py
--- a/jax/_src/interpreters/ad.py
+++ b/jax/_src/interpreters/ad.py
@@ -586,19 +586,19 @@ def zero_jvp(primitive, primals, tangents, **params):
 deflinear2(add_jaxvals_p, lambda t, *args: (t, t))
 
 def instantiate_zeros(tangent):
-  if type(tangent) is Zero:
-    return zeros_like_aval(tangent.aval)
-  else:
+  if type(tangent) is not Zero:
     return tangent
+  return instantiate_zeros_aval(tangent.aval, tangent)
 
 # This function seems similar to instantiate_zeros, but it is sometimes used
 # to instantiate zero abstract units with a different aval
 def instantiate_zeros_aval(aval, tangent):
-  if type(tangent) is Zero:
-    assert tangent.aval == aval
-    return zeros_like_aval(aval)
-  else:
+  if type(tangent) is not Zero:
     return tangent
+  assert tangent.aval == aval
+  if jax.dtypes.issubdtype(aval.dtype, jax.dtypes.extended):
+    return aval.dtype._rules.make_tangent(aval.shape, aval.dtype)
+  return zeros_like_aval(aval)
 
 @lu.transformation_with_aux
 def traceable(in_tree, *primals_and_tangents):
diff --git a/jax/_src/prng.py b/jax/_src/prng.py
--- a/jax/_src/prng.py
+++ b/jax/_src/prng.py
@@ -472,6 +472,23 @@ def full(shape, fill_value, dtype):
     # the outset.
     return random_wrap(key_data, impl=dtype._impl)
 
+  @staticmethod
+  def make_tangent(shape, dtype):
+    physical_shape = (*shape, *dtype._impl.key_shape)
+    def not_implemented(name):
+      def func(*args):
+        raise NotImplementedError(f"Cannot call {name} on tangent of PRNG key.")
+      return func
+    impl = PRNGImpl(
+      key_shape=dtype._impl.key_shape,
+      seed=not_implemented('seed'),
+      split=not_implemented('split'),
+      random_bits=not_implemented('random_bits'),
+      fold_in=not_implemented('fold_in'),
+      name=f"{dtype._impl.name}_tangent",
+      tag=f"{dtype._impl.tag}_t")
+    return random_wrap(jnp.zeros(physical_shape, dtype='uint32'), impl=impl)
+
   @staticmethod
   def physical_element_aval(dtype) -> core.ShapedArray:
     return core.ShapedArray(dtype._impl.key_shape, jnp.dtype('uint32'))
@@ -594,6 +611,20 @@ def device_put_replicated(val, aval, sharding, devices):
     return random_wrap(physical_result, impl=aval.dtype._impl)
 
 
+class KeyTangentTy(dtypes.ExtendedDType):
+  """A dtype to use for the tangent of a PRNGKey"""
+  _impl: PRNGImpl
+  type = dtypes.prng_key
+
+  @property
+  def _rules(self):
+    raise ValueError("Cannot perform operations on the tangent of a PRNGKey.")
+
+  @property
+  def name(self) -> str:
+    return f'key_tangent<{self._impl.tag}>'
+
+
 class KeyTy(dtypes.ExtendedDType):
   _impl: PRNGImpl  # TODO(mattjj,frostig): protocol really
   _rules = KeyTyRules


Test patch:
diff --git a/tests/random_test.py b/tests/random_test.py
--- a/tests/random_test.py
+++ b/tests/random_test.py
@@ -1107,6 +1107,36 @@ class A: pass
     with self.assertRaisesRegex(TypeError, 'unrecognized type .* PRNG'):
       jax.random.key(42, impl=A())
 
+  def test_keyarray_custom_vjp(self):
+    # Regression test for https://github.com/google/jax/issues/18442
+    @jax.custom_vjp
+    def f(_, state):
+      return state
+    def _f_fwd(_, state):
+      return state, None
+    def _f_bwd(_, state_bar):
+      assert state_bar[1].dtype.name == "key<fry_t>"  # key tangent type
+      return state_bar
+    f.defvjp(_f_fwd, _f_bwd)
+    state = (8.0, jax.random.key(123))
+    result = jax.grad(lambda theta: f(theta, state)[0])(3.0)
+    self.assertEqual(result, 1.0)
+
+  def test_keyarray_custom_vjp_symbolic_zeros(self):
+    @jax.custom_vjp
+    def f(_, state):
+      return state
+    def _f_fwd(_, state):
+      return tree_util.tree_map(lambda x: x.value, state), None
+    def _f_bwd(_, state_bar):
+      self.assertTrue(dtypes.issubdtype(state_bar[1].dtype, dtypes.prng_key))
+      self.assertIsInstance(state_bar[1], jax.custom_derivatives.SymbolicZero)
+      return state_bar
+    f.defvjp(_f_fwd, _f_bwd, symbolic_zeros=True)
+    state = (8.0, jax.random.key(123))
+    result = jax.grad(lambda theta: f(theta, state)[0])(3.0)
+    self.assertEqual(result, 1.0)
+
   # TODO(frostig,mattjj): more polymorphic primitives tests