optax.measure_with_ema

optax.measure_with_ema(measure: Callable[[base.Updates], base.ArrayTree], decay: jax.typing.ArrayLike, debias: bool = True, accumulator_dtype: Any | None = None) → base.GradientTransformationExtraArgs

Take a measurement and record it with exponential moving average.

Parameters:

• measure – User callable taking as inputs updates and returning desired measurement.

• decay – Decay rate for the exponential moving average.

• debias – Whether to debias the exponential moving average.

• accumulator_dtype – Optional dtype for the exponential moving average accumulator.

Returns:

A gradient transformation that captures measurements defined by the user, and records them with exponential moving average.

See also

optax.monitor()