from oryx.core.interpreters.harvest import sow, reap
import jax


@jax.jit
def f(x):
    sow(x, tag="d", name="x")

@jax.jit
def g(x):
    return reap(f, tag="d")(x=x)

g(5)  # {}