Skip to content

Sowlib scan - #5604

Open
samanklesaria wants to merge 5 commits into
google:mainfrom
samanklesaria:sowlib_scan
Open

samanklesaria wants to merge 5 commits into
google:mainfrom
samanklesaria:sowlib_scan

Conversation

@samanklesaria

Copy link
Copy Markdown
Collaborator

This PR builds on the jax-transform version of sow introduced in #5575 giving it support for scan. This is part two in a series of PRs that will eventually cover all of #5556

_run now looks each primitive up in _DISPATCH, falling back to binding it,
and threads a mtxrun          | unknown script 'mtx-context.lua' or 'mtx-mtx-context.lua' dict that dispatch rules can use for shared state —
currently the memo table for _contains_sow, which otherwise re-walks nested
jaxprs once per enclosing equation.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant