Adjustments generate_traj/_batch

If we declare mean and covariance as trajectories and try to tree_map with elite_u_trajs (batch) we get a "mismatch custom node data" because of batch_dim = None.