Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Stateful now supports vmap #466

Merged
merged 3 commits into from
Sep 12, 2023
Merged

Stateful now supports vmap #466

merged 3 commits into from
Sep 12, 2023

Commits on Sep 12, 2023

  1. Updated stateful operations to be able to support vmap'ing.

    This is actually a minor breaking change: the API has changed from
    StateIndex(callable_returning_value) to just StateIndex(value).
    
    This change also introduces substantially updated documentation on
    stateful layers, which should help a lot.
    patrick-kidger committed Sep 12, 2023
    Configuration menu
    Copy the full SHA
    2d4c553 View commit details
    Browse the repository at this point in the history
  2. Stateful operations now support creating states multiple times, such …

    …that they are compatible with the original model.
    patrick-kidger committed Sep 12, 2023
    Configuration menu
    Copy the full SHA
    61e47b8 View commit details
    Browse the repository at this point in the history
  3. Configuration menu
    Copy the full SHA
    3c8d4e8 View commit details
    Browse the repository at this point in the history