Jax a for loop that handles values with different shapes

[eadadi]

Hi all,

I have a traj dictionary that contains keys of different shape.

traj = {'action': ..., 'goal':..., etc}

I have a for loop that manipulates traj:

for k,v in traj.items():
  v = ...
  if k!='reward':
    v = ....
  traj[k] = v

Obviously, I want to parallelize this for loop.
I thought of using the scan method.
I did something like this:

def op(carry, k_v):
  k, v = k_v
  v = ..
  if k!=0: #I changed k to 0 if it is 'reward' and else to 1, to avoid jax error. not sure if necessary
    v = ..
  return carry + [(k,v)], None
inputs = jnp.array([[0,v] if k==''reward" else [1,v] for k,v in traj.items()])
results, _ = jax.lax.scan(operation, [], inputs)
traj = dict(results)

I currently get error over the jnp.array(..) line because the shape of the values are different.
And if I avoid using converting it to jnp.array, and instead do

inputs = [[0,v] if k==''reward" else [1,v] for k,v in traj.items()]

I instead get an error in the scan method that looks like:

ValueError: scan got value with no leading axis to scan over: 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1.

Can I get some help to overcome this?

[eadadi]

The way I solved this is by using:

jax.tree_util.tree_map(op, traj, keys)

where op is the logic I applied over any value, and keys is the keys of traj that I needed to use also inside that logic