wax.universal.eager_ewma

Universal Exponential moving average module and unroll implementations.

Functions

dynamic_unroll_fori_loop_universal(fun, …)	Unroll a TransformedWithState function using fori_loop in an universal way.
dynamic_unroll_tf(fun, params, state, rng, …)	Unroll a TransformedWithState function using tensorflow scan function.
dynamic_unroll_universal(fun, params, state, rng)	Unroll a TransformedWithState function using jax.lax.scan.
static_unroll_universal(fun, params, state, rng)	Unroll a TransformedWithState function using static_scan implemented in an universal way.

Classes

EagerEWMA(*args, **kwargs)

Universal Exponential moving average module.