kd.train.RngStream

class kauldron.train.RngStream(

name: str,

*,

init: bool = True,

train: bool = True,

eval: bool = False,

per_step: bool = True,

)

Bases: object

Info on one rng stream.

See doc at https://kauldron.rtfd.io/en/latest/eval.html#rng-streams

name

Stream name

Type:

str

init

Whether the rng is used in model.init

Type:

bool

train

Whether the rng is used in train (is_training=True)

Type:

bool

eval

Whether the rng is used in eval (is_training=False)

Type:

bool

per_step

Whether the rng is different at each step

Type:

bool

name: str

init: bool = True

train: bool = True

eval: bool = False

per_step: bool = True

make(

rng: kauldron.random.random.PRNGKey,

*,

step: int | None = None,

key: str | None = None,

) → kauldron.random.random.PRNGKey

Create the rng from the global root rng.

Parameters:

• rng – The root rng, common to all processes

• step – Current model step

• key – Additional string (e.g. train, init,…) to fold in

Returns:

The new rng