jax
5f1f29e4 - fix jaxpr invar avals

Commit
5 years ago
fix jaxpr invar avals When an input argument is known to partial_eval, the corresponding jaxpr (which represents only the staged-out, unknown part of the computation) has a dummy placeholder input with type (i.e. aval) abstract_unit. However, the avals attached to the invars were recording the aval of the known value in that case. cf. #2299
Author
Parents
Loading