We could simplify even more by deriving the torch prng key from the step counter and using heavyball'sdeterministic precond. However, I don't know if that'd take too far.

The docs for this PR livehere. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

@ClashLuke thanks, much better way to do this without worrying about restoring a device specific torch rng state.. merged with small fix