flax
Update libraries to use JAX's limited (and ill-advised) trace-state-querying APIs rather than depending on JAX's deeper internals, which are about to change.
#4219

Open

Update libraries to use JAX's limited (and ill-advised) trace-state-querying APIs rather than depending on JAX's deeper internals, which are about to change. #4219

copybara-service wants to merge 1 commit into main from test_677843398

copybara-service force pushed 1 year ago

copybara-service changed the title ~~Hide JAX's internal tracing state and update libraries to use limited trace-state-querying APIs as needed. This is prep work for stackless which will change those internals while preserving the API.~~ Update libraries to use JAX's limited (and ill-advised) trace-state-querying APIs rather than depending on JAX's deeper internals, which are about to change. 1 year ago

Update libraries to use JAX's limited (and ill-advised) trace-state-q…

815f9f53

copybara-service force pushed to 815f9f53 1 year ago

Reviewers

No reviews

Assignees

No one assigned

Labels

None yet

Milestone

No milestone

flax Update libraries to use JAX's limited (and ill-advised) trace-state-querying APIs rather than depending on JAX's deeper internals, which are about to change. #4219 Open

Update libraries to use JAX's limited (and ill-advised) trace-state-querying APIs rather than depending on JAX's deeper internals, which are about to change. #4219

flax
Update libraries to use JAX's limited (and ill-advised) trace-state-querying APIs rather than depending on JAX's deeper internals, which are about to change.
#4219

Open