Computes a moving average of `x`. The moving average updates will be set in OutputCollection.state_updates. Args: x: A Tensor of shape cfg.shape. Returns: Returns the current moving average.
(self, x: Tensor)
| 1564 | """Computes a moving average of `x`. |
| 1565 | |
| 1566 | The moving average updates will be set in OutputCollection.state_updates. |
| 1567 | |
| 1568 | Args: |
| 1569 | x: A Tensor of shape cfg.shape. |
| 1570 | |
| 1571 | Returns: |
| 1572 | Returns the current moving average. |
| 1573 | """ |
| 1574 | cfg = self.config |
| 1575 | weight = jnp.maximum(cfg.min_weight, 1.0 / (1 + self.parameters["count"])) |
| 1576 | new_moving_average = (1 - weight) * self.parameters["value"] + weight * x |
| 1577 | self.add_state_update("value", new_moving_average) |
| 1578 | self.add_state_update("count", 1 + self.parameters["count"]) |
| 1579 | return new_moving_average |
nothing calls this directly
no test coverage detected