pypomp.LearningRate.to_array

LearningRate.to_array(param_names: list[str], M: int) → Array

Convert the learning rates into a JAX array.

Parameters:

• param_names (list of str) – Parameter names in canonical order.

• M (int) – Number of iterations in the training schedule.

Returns:

A 2D array of shape (M, n_params) where each column is the learning rate schedule for a parameter.

Return type:

jax.Array