Debug array in jax vmap function

Dear jax experts I need your kind help.

Here is a working example (I have follow the advise to simplify my code, although I am not an expert on jax neither on Python to guess what is the heart of the mechanism involved in vmap)

def jax_kernel(rng_key, logpdf, position, log_prob):
    key, subkey = jax.random.split(rng_key)
    move_proposals = jax.random.normal(key, shape=position.shape)* 0.1   
    proposal = position + move_proposals
    proposal_log_prob = logpdf(proposal)
    return proposal, proposal_log_prob

def jax_sampler(rng_key, n_samples, logpdf, initial_position):
    
    def mh_update(i, state):
        key, positions, log_prob = state
        _, key = jax.random.split(key)        
        print(f"mh_update: positions[{i-1}]:",jnp.asarray(positions[i-1]))
        new_position, new_log_prob = jax_kernel(key,logpdf,positions[i-1],log_prob)
            
        positions=positions.at[i].set(new_position)
        return (key, positions, new_log_prob)
    
    # all positions structure should be set before lax.fori_loop
    print("initial_position shape:",initial_position.shape)       
    all_positions = jnp.zeros((n_samples,)+initial_position.shape)  
    all_positions=all_positions.at[0,0].set(1.)
    all_positions=all_positions.at[0,1].set(2.)
    all_positions=all_positions.at[0,2].set(2.)
    print("all_positions init:",all_positions.shape)
    logp = logpdf(all_positions[0])
    
    # use of a for-loop to be able to debug mh_update instead of a jax.fori_loop
    initial_state = (rng_key,all_positions, logp)
    val = initial_state
    for i in range(1, n_samples):
        val = mh_update(i, val)
    rng_key, all_positions, log_prob = val
    # return all the positions of the parameters (n_chains, n_samples, n_dim)
    return all_positions

def func(par):
    xi = jnp.asarray(sci_stats.uniform.rvs(size=10))
    val = xi*par[1]+par[0]
    return jnp.sum(jax.scipy.stats.norm.logpdf(x=val,loc=yi,scale=par[2]))
    

n_dim = 3          # number of parameters ie. (a,b,s)
n_samples = 5      # number of samples per chain
n_chains = 4       # number of MCMC chains
rng_key = jax.random.PRNGKey(42)
rng_keys = jax.random.split(rng_key, n_chains) 
initial_position = jnp.ones((n_dim, n_chains))                      
print("main initial_position shape",initial_position.shape)
run = jax.vmap(jax_sampler, in_axes=(0, None, None, 1), out_axes=0) 
all_positions = run(rng_keys,n_samples,lambda p: func(p),initial_position)
print("all_positions:",all_positions)

Then my question concerns the dimension evolution print(f"mh_update: positions[{i-1}]:",jnp.asarray(positions[i-1])). I do not understand why positions[i-1]starts with dimension n_dim and then switches to n_chains x n_dim?

Thanks in advance for your comments?

Here is the complete output:

main initial_position shape (3, 4)
initial_position shape: (3,)
all_positions init: (5, 3)
mh_update: positions[0]: [1. 2. 2.]
mh_update: positions[1]: Traced<ShapedArray(float32[3])>with<BatchTrace(level=1/0)>
  with val = DeviceArray([[0.9354116 , 1.7876872 , 1.8443539 ],
                          [0.9844745 , 2.073029  , 1.9511036 ],
                          [0.98202926, 2.0109322 , 2.094176  ],
                          [0.9536771 , 1.9731759 , 2.093319  ]], dtype=float32)
       batch_dim = 0
mh_update: positions[2]: Traced<ShapedArray(float32[3])>with<BatchTrace(level=1/0)>
  with val = DeviceArray([[1.0606856, 1.6707807, 1.8377957],
                          [1.0465866, 1.9754674, 1.7009288],
                          [1.1107644, 2.0142047, 2.190575 ],
                          [1.0089972, 1.9953227, 1.996874 ]], dtype=float32)
       batch_dim = 0
mh_update: positions[3]: Traced<ShapedArray(float32[3])>with<BatchTrace(level=1/0)>
  with val = DeviceArray([[1.0731456, 1.644405 , 2.1343162],
                          [1.0599504, 2.0121546, 1.6867112],
                          [1.0585173, 1.9661485, 2.1573594],
                          [1.1213307, 1.9335203, 1.9683584]], dtype=float32)
       batch_dim = 0
all_positions: [[[1.         2.         2.        ]
  [0.9354116  1.7876872  1.8443539 ]
  [1.0606856  1.6707807  1.8377957 ]
  [1.0731456  1.644405   2.1343162 ]
  [1.0921828  1.5742197  2.058759  ]]

 [[1.         2.         2.        ]
  [0.9844745  2.073029   1.9511036 ]
  [1.0465866  1.9754674  1.7009288 ]
  [1.0599504  2.0121546  1.6867112 ]
  [1.0835105  2.0051234  1.4766487 ]]

 [[1.         2.         2.        ]
  [0.98202926 2.0109322  2.094176  ]
  [1.1107644  2.0142047  2.190575  ]
  [1.0585173  1.9661485  2.1573594 ]
  [1.1728328  1.981367   2.180744  ]]

 [[1.         2.         2.        ]
  [0.9536771  1.9731759  2.093319  ]
  [1.0089972  1.9953227  1.996874  ]
  [1.1213307  1.9335203  1.9683584 ]
  [1.1148386  1.9598911  2.1721165 ]]]

from Debug array in jax vmap function