OptimHistory.init_position_history()

OptimHistory.init_position_history()#

static OptimHistory.init_position_history(position, epochs)[source]#

Allocates zero-filled arrays for a position history.

The returned arrays have one extra leading epoch dimension.

Return type:

Position (dict[str, Any])

Examples

>>> import jax.numpy as jnp
>>> from liesel.optim import OptimHistory
>>> from liesel.optim.types import Position
>>> position = Position({"theta": jnp.ones((2,))})
>>> history = OptimHistory.init_position_history(position, epochs=3)
>>> history["theta"].shape
(3, 2)
>>> history["theta"].sum()
Array(0., dtype=float32)