Stackless, redux #38
Stackless, redux #38
Conversation
`pytest tests` passes with these changes
Signed-off-by: Nathaniel Starkman <nstarman@users.noreply.github.com>
nstarman
force-pushed
the
stackless-redux
branch
5 times, most recently
from
4fbc627 to
c0c7530
Compare
Signed-off-by: Nathaniel Starkman <nstarman@users.noreply.github.com>
Signed-off-by: Nathaniel Starkman <nstarman@users.noreply.github.com>
nstarman
force-pushed
the
stackless-redux
branch
from
8d3599d to
4e79652
Compare
|
Nice! Thank you for putting this together. Do you have an idea how to tackle #37 (comment) ? |
| dynamic = jtu.tree_map( | ||
| ft.partial(_wrap_tracer, trace), | ||
| dynamic, | ||
| is_leaf=_is_value, | ||
| ) | ||
| fn, args, kwargs = eqx.combine(dynamic, static) | ||
| out = fn(*args, **kwargs) | ||
| with core.set_current_trace(trace): |
There was a problem hiding this comment.
@patrick-kidger @mattjj does this look like the place to check that Quax values are wrapped into tracers?
#37 (comment)
There was a problem hiding this comment.
I can do some testing, but an informed starting point would be great!
There doesn't appear to be an obvious place in _QuaxTrace anymore, so I was looking at here, but it's also not obviously the right place for checking the tracer-wrapping.
|
@patrick-kidger |
|
Ah, interesting. In that case I think I'm happy to merge this as-is. This is all still pretty new so we'll just have to see how it goes, I suppose! If you'd like to bump the version number then I'm happy to do a new release. :) Thank you for taking over this PR! |
I'll do that! |
Hi @patrick-kidger @mattjj!
I cherry-picked @mattjj's commits, fixed the CI, and made the suggested refactor.
I'm happy to push these fixes as a PR to @mattjj's branch instead of this PR.
The changes in jax also require python 3.10+