Would using lists rather than jax.numpy arrays lead to more accurate numerical transformations?

[YousefElbrolosy]

Hi, I am doing a project with RNNs using jax and flax and I have noticed some behavior that I do not really understand.

My code is basically an optimization loop where the user provides the initial parameters for the system they want to optimize. This system is divided onto several time steps. He feeds the initial input into the first time step of the the system, gets a certain output, feeds this output into a RNN which returns the parameters for the following time step and so on. Then it is optimized using adam (particularly using optax).

Now the user inputs his initial paramaters as a dict and then there is a function called prepare_parameters_from_dict that basically converts this dict into a list of lists ( or a list of jnp arrays for that matter ).

My question/observation is when I make this function return a list of jnp.arrays instead of a list of lists, the property I am optimizing is an order of magnitude worse!

For example, using a list of lists outputs 0.9997 and a list of jnp.arrays outputs 0.998 (the closer to one the better).

Noting: the RNN output a list of jnp.arrays (it is using flax linnen) and everything in the code remains the same.

Here are said function:

Outputing list of lists:

def prepare_parameters_from_dict(params_dict):
    """
    Convert a nested dictionary of parameters to a flat list and record shapes.

    Args:
        params_dict: Nested dictionary of parameters.

    Returns:
        tuple: Flattened parameters list and list of shapes.
    """
    res = []
    shapes = []
    for value in params_dict.values():
        flat_params = jax.tree_util.tree_leaves(value)
        res.append(flat_params)
        shapes.append(len(flat_params))
    return res, shapes

Using list of jnp.arrays:

def prepare_parameters_from_dict(params_dict):
    """
    Convert a nested dictionary of parameters to a flat list and record shapes.

    Args:
        params_dict: Nested dictionary of parameters.

    Returns:
        tuple: Flattened parameters list and list of shapes.
    """
    res = []
    shapes = []
    for value in params_dict.values():
        flat_params = jax.tree_util.tree_leaves(value)
        res.append(jnp.array(flat_params))
        shapes.append(jnp.array(flat_params).shape[0])
    return res, shapes

and this is an example of the users input initial params:

initial_params = {
    "param1": {
        "gamma": 0.1,
        "delta": -3 * jnp.pi / 2,
    }
}

The rest of the code remains exactly the same for both.

After optimization if for example there were five time steps, this is how the final optimized params for each time step would look like:

using list of jnp.arrays:

[[Array([ 0.1       , -4.71238898], dtype=float64)],
 [Array([-0.97106537, -0.03807388], dtype=float64)],
 [Array([-1.17050792, -0.01463591], dtype=float64)],
 [Array([-0.77229875, -0.0124556 ], dtype=float64)],
 [Array([-1.56113376, -0.01103598], dtype=float64)]]

using list of lists:

[[ [0.1       , -4.71238898] ]],
 [Array([-0.97106537, -0.03807388], dtype=float64)],
 [Array([-1.17050792, -0.01463591], dtype=float64)],
 [Array([-0.77229875, -0.0124556 ], dtype=float64)],
 [Array([-1.56113376, -0.01103598], dtype=float64)]]

Would such a difference in behavior be due to how jax handles grad and jit and others with lists compared to jnp.arrays or am I missing something?

[YousefElbrolosy]

For those who had a similar Question/Problem refer to this stackoverflow question, I asked the same one there and got an answer.

[jakevdp]

sorry, I somehow missed the question being asked here!

[YousefElbrolosy]

All good! Thanks for your answer over there!