Implement process_call for LinearizeTrace
#25481
Conversation
jax/_src/interpreters/ad.py
Outdated
| out_tangents = map(tangent_trace.to_jaxpr_tracer, out_tangents) | ||
| jaxpr, consts, attrs_tracked = tangent_trace.to_jaxpr(out_tangents) | ||
| num_residuals = len(consts) | ||
| del attrs_tracked # TODO: attrs |
There was a problem hiding this comment.
Raise an exception if there are any? (You can also unpack into an empty tuple to check the same, though with no explanation of what the error means.)
jax/_src/lax/lax.py
Outdated
| @@ -4153,9 +4153,10 @@ def _broadcast_in_dim_typecheck_rule( | |||
|
|
|||
| def _broadcast_in_dim_transpose_rule(ct, operand, *dyn_shape, | |||
| shape, broadcast_dimensions, sharding): | |||
| aval = core.get_aval(operand) | |||
There was a problem hiding this comment.
Huh, I would've thought operand would be an UndefinedPrimal here. Does get_aval work on that? Why this change?
There was a problem hiding this comment.
Good point. I changed it because operand was sometimes a literal. I thought the implementation was just missing that case (sometimes we assume things are tracers when they can be literals or tracers). But you're right, something else must be going wrong. I'll investigate.
There was a problem hiding this comment.
Figured it out. I was instantiating zeros inside the linear jaxpr rather than outside it. Gotta dedent your instantiate_zeros.
google-ml-butler
bot
added
kokoro:force-run
pull ready
dougalm
force-pushed
the
custom-linearize-process-call
branch
2 times, most recently
from
19058ba to
032799f
Compare
dougalm
force-pushed
the
custom-linearize-process-call
branch
from
032799f to
dea51cb
Compare
No description provided.