State for the Lion algorithm.
| 1653 | |
| 1654 | |
| 1655 | class ScaleByLionState(NamedTuple): |
| 1656 | """State for the Lion algorithm.""" |
| 1657 | |
| 1658 | count: chex.Array # shape=(), dtype=jnp.int32. |
| 1659 | mu: optax.Updates |
| 1660 | |
| 1661 | |
| 1662 | def scale_by_lion( |
no outgoing calls
no test coverage detected