A struct-of-arrays to an array-of-structs?

[dui1234]

Hi,

If I have a struct-of-arrays type of data and I want to convert it to an array-of-structs, how can I do that?

For example conversion from the params_stacked back to params in the code below:

import jax
import jax.numpy as jnp
from jax.nn import initializers
from jax.example_libraries import stax, optimizers

def get_NN(D_in,H,D_out,seed=5407):
    key = jax.random.PRNGKey(seed)
    key, sk1,sk2 = jax.random.split(key, 3)
    Winit = initializers.lecun_normal(dtype=jnp.float64)
    binit = initializers.normal(stddev=1.0, dtype=jnp.float64)
        
    net_init, net_forward_fn = stax.serial(stax.Dense(H, Winit, binit), stax.Relu,
                                stax.Dense(H, Winit, binit), stax.Relu,
                                stax.Dense(D_out, Winit, binit))
    output_shape, net_params = net_init(sk2, (D_in,))
        
    return [net_params,net_forward_fn]

def get_agent(s_params):
    policy_NN = get_NN(s_params['obs_dim'],s_params['H'],s_params['action_dim'])
    
    lr = 1e-3
    opt_init, opt_update, get_params = optimizers.adam(step_size=lr)
    opt_state = opt_init(policy_NN[0])
    
    return opt_state,policy_NN[1], opt_update, get_params

system_params = {'epsilon': 1.0,'obs_dim':4,'action_dim':2,'H':5,'N_swim':2}

NNs = []
for i in range(system_params['N_swim']):
    NN,forward_f,opt_update,get_params = get_agent(system_params)
    NNs.append(NN)

params = [x for x in NNs]
params_stacked = jax.tree_map(lambda *x: jnp.stack(x), *params)
#How about from params_stacked --> params??

[dui1234]

Sorry, I just found out that I can simply use:
params = [ jax.tree_map(lambda x: x[i], params_stacked) for i in range(system_params['N_swim']) ]

[mattjj]

It's a great question!

Indeed, as you found we tend to just offer APIs for working with the struct-of-arrays form. But we do have a prototype (untouched since the holiday break a year ago...) of providing an array-of-structures API, with automatic conversion to structure-of-arrays under the hood. See the tests here if you're curious. It was pretty lightweight; maybe someday we'll land it!