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