ott.solvers.nn.models.ModelBase.make_rng

ModelBase.make_rng(name)

Returns a new RNG key from a given RNG sequence for this Module.

The new RNG key is split from the previous one. Thus, every call to make_rng returns a new RNG key, while still guaranteeing full reproducibility.

TODO: Link to Flax RNG design note.

Parameters:

name (str) – The RNG sequence name.

Return type:

Any

Returns:

The newly generated RNG key.