Class _AdastarPerParamState

axlearn/common/optimizers.py:1859–1862 · view source on GitHub ↗

Source from the content-addressed store, hash-verified

1857	"""
1858
1859	class _AdastarPerParamState(flax_struct.PyTreeNode):
1860	gradient_ema: Optional[Tensor]
1861	gradient_square_ema: Tensor
1862	update_ema: Optional[Tensor]
1863
1864	class _AdastarState(flax_struct.PyTreeNode):
1865	count: Tensor

_initFunction · 0.85

_raw_updatesFunction · 0.85

_smoothed_updatesFunction · 0.85

_partitionFunction · 0.85

no outgoing calls

no test coverage detected