Extend beyond 1D

ziatdinovmax · Jan 2, 2024 · 3121fdb · 3121fdb
1 parent 7035dcc
commit 3121fdb
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/gpax/acquisition/optimize.py b/gpax/acquisition/optimize.py
@@ -58,7 +58,9 @@ def optimize_acq(rng_key: jnp.ndarray,
         ) from e
 
     def acq(x):
-        obj = -acq_fn(rng_key, model, jnp.array([x])[None], **kwargs)
+        x = jnp.array([x])
+        x = x[None] if x.ndim == 0 else x
+        obj = -acq_fn(rng_key, model, x, **kwargs)
         return jnp.reshape(obj, ())
 
     lower_bound = ensure_array(lower_bound)