jax.grad with def async functions?

[tylerflex]

I want to use jax.grad to differentiate the result of an async function. Here is a minimal example.

import jax

async def f(x):
    return x + 1

y = await f(1.0)
df = jax.grad(f)
dy = await df(1.0)

I get the following

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Cell In[4], line 10
      7 print(y)
      9 df = jax.grad(f)
---> 10 dy = await df(1.0)
     11 print(dy)

    [... skipping hidden 15 frame]

File ~/.pyenv/versions/3.10.9/lib/python3.10/site-packages/jax/_src/core.py:1322, in concrete_aval(x)
   1320 if hasattr(x, '__jax_array__'):
   1321   return concrete_aval(x.__jax_array__())
-> 1322 raise TypeError(f"Value {repr(x)} with type {type(x)} is not a valid JAX "
   1323                  "type")

TypeError: Value <coroutine object f at 0x10c091930> with type <class 'coroutine'> is not a valid JAX type

jax version '0.4.4'

can someone point me in the right direction? Is this even possible with jax?
Note: in my actual application, I have implemented f.def_vjp() in terms of two async functions themselves and get the same error.

Thank you!

[tylerflex]

Answered my own question finally, the key was to pass

f_async = lambda x: asyncio.run(f(x))

to jax.grad.