diff --git a/.bumpversion.toml b/.bumpversion.toml
index a468f23..4ded5eb 100644
--- a/.bumpversion.toml
+++ b/.bumpversion.toml
@@ -1,5 +1,5 @@
 [tool.bumpversion]
-current_version = "v0.5.1"
+current_version = "v0.5.2"
 commit = true
 commit_args = "--no-verify"
 tag = true
diff --git a/README.md b/README.md
index 3b60a08..5c0ad80 100644
--- a/README.md
+++ b/README.md
@@ -1,5 +1,5 @@
 # invrs-opt - Optimization algorithms for inverse design
-`v0.5.1`
+`v0.5.2`
 
 ## Overview
 
diff --git a/pyproject.toml b/pyproject.toml
index ffd0424..bd55b78 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,7 +1,7 @@
 [project]
 
 name = "invrs_opt"
-version = "v0.5.1"
+version = "v0.5.2"
 description = "Algorithms for inverse design"
 keywords = ["topology", "optimization", "jax", "inverse design"]
 readme = "README.md"
diff --git a/src/invrs_opt/__init__.py b/src/invrs_opt/__init__.py
index 46687db..5f0cc44 100644
--- a/src/invrs_opt/__init__.py
+++ b/src/invrs_opt/__init__.py
@@ -3,7 +3,7 @@
 Copyright (c) 2023 The INVRS-IO authors.
 """
 
-__version__ = "v0.5.1"
+__version__ = "v0.5.2"
 __author__ = "Martin F. Schubert <mfschubert@gmail.com>"
 
 from invrs_opt.lbfgsb.lbfgsb import density_lbfgsb as density_lbfgsb
diff --git a/src/invrs_opt/lbfgsb/lbfgsb.py b/src/invrs_opt/lbfgsb/lbfgsb.py
index f29bddc..871ccd1 100644
--- a/src/invrs_opt/lbfgsb/lbfgsb.py
+++ b/src/invrs_opt/lbfgsb/lbfgsb.py
@@ -258,7 +258,7 @@ def _init_pure(params: PyTree) -> Tuple[PyTree, JaxLbfgsbDict]:
         (
             latent_params,
             jax_lbfgsb_state,
-        ) = jax.pure_callback(  # type: ignore[attr-defined]
+        ) = jax.pure_callback(
             _init_pure,
             _example_state(params, maxcor),
             initialize_latent_fn(params),
@@ -304,7 +304,7 @@ def _update_pure(
         (
             flat_latent_params,
             jax_lbfgsb_state,
-        ) = jax.pure_callback(  # type: ignore[attr-defined]
+        ) = jax.pure_callback(
             _update_pure,
             (flat_latent_grad, jax_lbfgsb_state),
             flat_latent_grad,
@@ -542,19 +542,19 @@ def from_jax(cls, state_dict: Dict[str, jnp.ndarray]) -> "ScipyLbfgsbState":
         """Converts a dictionary of jax arrays to a `ScipyLbfgsbState`."""
         state_dict = copy.deepcopy(state_dict)
         return ScipyLbfgsbState(
-            x=onp.asarray(state_dict["x"], dtype=onp.float64),
+            x=onp.array(state_dict["x"], dtype=onp.float64),
             converged=onp.asarray(state_dict["converged"], dtype=bool),
             _maxcor=int(state_dict["_maxcor"]),
             _line_search_max_steps=int(state_dict["_line_search_max_steps"]),
             _ftol=onp.asarray(state_dict["_ftol"], dtype=onp.float64),
             _gtol=onp.asarray(state_dict["_gtol"], dtype=onp.float64),
-            _wa=onp.asarray(state_dict["_wa"], onp.float64),
-            _iwa=onp.asarray(state_dict["_iwa"], dtype=FORTRAN_INT),
+            _wa=onp.array(state_dict["_wa"], onp.float64),
+            _iwa=onp.array(state_dict["_iwa"], dtype=FORTRAN_INT),
             _task=_s60_str_from_array(state_dict["_task"]),
             _csave=_s60_str_from_array(state_dict["_csave"]),
-            _lsave=onp.asarray(state_dict["_lsave"], dtype=FORTRAN_INT),
-            _isave=onp.asarray(state_dict["_isave"], dtype=FORTRAN_INT),
-            _dsave=onp.asarray(state_dict["_dsave"], dtype=onp.float64),
+            _lsave=onp.array(state_dict["_lsave"], dtype=FORTRAN_INT),
+            _isave=onp.array(state_dict["_isave"], dtype=FORTRAN_INT),
+            _dsave=onp.array(state_dict["_dsave"], dtype=onp.float64),
             _lower_bound=onp.asarray(state_dict["_lower_bound"], dtype=onp.float64),
             _upper_bound=onp.asarray(state_dict["_upper_bound"], dtype=onp.float64),
             _bound_type=onp.asarray(state_dict["_bound_type"], dtype=int),