ott.neural.networks.icnn.ICNN.make_rng

Contents

ott.neural.networks.icnn.ICNN.make_rng#

ICNN.make_rng(name='params')#

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.

NOTE: if an invalid name is passed (i.e. no RNG key was passed by the user in .init or .apply for this name), then name will default to 'params'.

TODO: Link to Flax RNG design note.

Parameters:

name (str) – The RNG sequence name.

Return type:

Array

Returns:

The newly generated RNG key.