Method _create_layer_parameter_specs

axlearn/common/layers.py:1547–1564

(self)

        cfg = self.config
        return {
            "count": ParameterSpec(
                shape=[],
                dtype=jnp.int32,
                mesh_axes=(None,),
                initializer=constant_initializer(0),
                weight_decay_scale=0,
            ),
            "value": ParameterSpec(
                shape=cfg.shape,
                dtype=jnp.float32,
                mesh_axes=(None,),
                initializer=constant_initializer(0.0),
                weight_decay_scale=0,
            ),
        }

    def forward(self, x: Tensor) -> Tensor:
        """Computes a moving average of `x`.

        The moving average updates will be set in OutputCollection.state_updates.