Const sows dropped in transformations

[jkramar]

Here's an example:

def f(x):
  harvest.sow(0, name='x', tag='tag')
  return x


print(harvest.reap(f, tag='tag')(1))
print(harvest.reap(lambda x: jax.remat(f), tag='tag')(1))
print(harvest.reap(lambda x: jax.jit(f), tag='tag')(1))
print(harvest.reap(lambda x: jax.lax.cond(x == 1, f, f, x), tag='tag')(1))
print(
    harvest.reap(
        lambda x: jax.lax.fori_loop(0, 1, lambda i, xx: f(xx), x), tag='tag'
    )(1)
)
print(
    harvest.reap(
        lambda x: jax.lax.fori_loop(0, x, lambda i, xx: f(xx), x), tag='tag'
    )(1)
)

The first of these reaps contains 'x': 0, but the others don't.