[Python] Would using lists rather than jax.numpy arrays lead to more accurate numerical...

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.
"""
flat_params = []
param_shapes = []

# returns a flat list of the leaves
def flatten_dict(d):
result = []
for key, value in d.items():
if isinstance(value, dict):
result.extend(flatten_dict(value))
else:
result.append(value)
return result

# flatten each top-level gate
for gate_name, gate_params in params_dict.items():
if isinstance(gate_params, dict):
# Extract parameters for this gate
gate_flat_params = flatten_dict(gate_params)
else:
# If already a flat array
gate_flat_params = gate_params
# this is checking if use can enter sth like {'gate1': 1} instead of {"gate1": {"param1": 1}}
if not (isinstance(gate_flat_params, list)):
flat_params.append([gate_flat_params])
param_shapes.append(1)
else:
flat_params.append(gate_flat_params)
param_shapes.append(len(gate_flat_params))

return flat_params, param_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.
"""
flat_params = []
param_shapes = []

# returns a flat list of the leaves
def flatten_dict(d):
result = []
for key, value in d.items():
if isinstance(value, dict):
result.extend(flatten_dict(value))
else:
result.append(value)
return result

# flatten each top-level gate
for gate_name, gate_params in params_dict.items():
if isinstance(gate_params, dict):
# Extract parameters for this gate
gate_flat_params = jnp.array(flatten_dict(gate_params))
else:
# If already a flat array
gate_flat_params = jnp.array(gate_params)
# this is checking if use can enter sth like {'gate1': 1} instead of {"gate1": {"param1": 1}}

flat_params.append(gate_flat_params)
param_shapes.append(gate_flat_params.shape)

return flat_params, param_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?

Continue reading...