make initial guess also vectorized

Author: edgarWolf

driftbench/data_generation/solvers.py

@@ -59,7 +59,7 @@ def __init__(self, f, w0, max_fit_attemps, vectorize=True):
             self.min_func = jit(
                 vmap(
                     partial(_minimize),
-                    in_axes=(None, None, None, None, 0, 0, 0, 0, 0, 0),
+                    in_axes=(None, None, None, 0, 0, 0, 0, 0, 0, 0),
                 ),
                 static_argnums=(0, 1, 2),
             )
@@ -98,7 +98,7 @@ def _solve_sequentially(self, X, callback=None):
         return jnp.array(coefficients)
 
     def _solve_vectorized(self, X):
-        solution = self.w0
+        solution = jnp.tile(self.w0, (len(X), 1))
         l_x0_mat, l_x1_mat, l_x2_mat, l_y0_mat, l_y1_mat, l_y2_mat = (
             self._latents_to_array(X)
         )