QuantityGenerator.generate()

QuantityGenerator.generate(prng_key, model_state, epoch)

Generates a new quantity based on the model and PRNG state.

Parameters

• prng_key (Any) – The key for JAX’ pseudo-random number generator.

• model_state (Any) – Current model state.

• epoch (EpochState) – Current epoch state.

Return type

TypeVar(TGeneratedQuantity, bound= GeneratedQuantity, covariant=True)