flax
Make `TrainState`'s `step` possibly jax.Array. This makes `replicate` valid for type checking.
#3763

Merged

Make `TrainState`'s `step` possibly jax.Array. This makes `replicate` valid for type checking. #3763

ghost merged 1 commit into main from test_615941638

copybara-service force pushed 1 year ago

copybara-service force pushed 1 year ago

Make `TrainState`'s `step` possibly jax.Array. This makes `replicate`…

718aa8cc

copybara-service force pushed to 718aa8cc 1 year ago

copybara-service closed this 1 year ago

copybara-service deleted the test_615941638 branch 1 year ago

Reviewers

No reviews

Assignees

No one assigned

Labels

None yet

Milestone

No milestone