Example:
def body(carry, _):
return harvest.sow(carry, name='x', mode='clobber', tag='tag'), None
f = jax.jit(lambda x: jax.lax.scan(body, x, None, 5))
harvest.harvest(f, tag='tag')({'x': 1}, 1)
This raises ValueError: safe_zip() argument 2 is longer than argument 1.
Interestingly, it does work without the planted value, or without the jit, or if we replace harvest.harvest with harvest.plant.
Example:
This raises
ValueError: safe_zip() argument 2 is longer than argument 1.Interestingly, it does work without the planted value, or without the jit, or if we replace
harvest.harvestwithharvest.plant.