Skip to content

Give RngStream the same interface as Jax's stateful RNGs #5046

@samanklesaria

Description

@samanklesaria

Jax will likely support a stateful PRNG object similar to Flax's RngStreams in the future: jax-ml/jax#28845.

Eventually, we'll want to have flax use these native flax objects. To ease that eventual transition, we should give RngStreams the same interface (specifically, the key and split methods) and express the new explicit-rng docs in terms of this interface.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions