[scan] don't hoist loop-invariant computations in scan, just forward #28993
Conversation
| in_fwd: list[int | None] = pe._jaxpr_forwarding(jaxpr_known.jaxpr) | ||
| in_fwd = [f if out_idx >= num_knowns_out and f is not None and | ||
| (f < num_consts_known or f >= num_consts_known + num_carry_known) | ||
| and isinstance(known_inputs[f], jax.Array) # no np.ndarrays |
There was a problem hiding this comment.
In the other forwarding logic, we use the _maybe_put helper to (I think?) support forwarding numpy arrays, but casting them to CPU device arrays. It seems like that should work here too, but I could well be missing something. What's the motivation for one or the other?
There was a problem hiding this comment.
I added this "dont forward numpy.ndarrays" line because this PR caused a failure in some Pathways test without it. I think _maybe_put was perhaps hardcoding the Arrays to be on host, rather than on default device, and that caused a really low-level error.
Out of laziness I decided to try just skipping the forwarding for numpy.ndarrays because I suspect we don't really need it. Another fix might be to improve _maybe_put, but I haven't looked into that at all.
mattjj
force-pushed
the
scan-dont-hoist
branch
2 times, most recently
from
6e08f0d to
08bc676
Compare
No description provided.