QuantityGenerator.generate()#

QuantityGenerator.generate(prng_key, model_state, epoch)[source]#

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)