invrs-opt 0.4.0__py3-none-any.whl → 0.10.3__py3-none-any.whl

invrs_opt/__init__.py

@@ -3,8 +3,19 @@
 Copyright (c) 2023 The INVRS-IO authors.
 """
 
-__version__ = "v0.4.0"
+__version__ = "v0.10.3"
 __author__ = "Martin F. Schubert <mfschubert@gmail.com>"
 
-from invrs_opt.lbfgsb.lbfgsb import density_lbfgsb as density_lbfgsb
-from invrs_opt.lbfgsb.lbfgsb import lbfgsb as lbfgsb
+from invrs_opt import parameterization as parameterization
+
+from invrs_opt.optimizers.lbfgsb import (
+    density_lbfgsb as density_lbfgsb,
+    lbfgsb as lbfgsb,
+    levelset_lbfgsb as levelset_lbfgsb,
+)
+
+from invrs_opt.optimizers.wrapped_optax import (
+    density_wrapped_optax as density_wrapped_optax,
+    levelset_wrapped_optax as levelset_wrapped_optax,
+    wrapped_optax as wrapped_optax,
+)

invrs_opt/experimental/client.py

@@ -4,15 +4,14 @@ Copyright (c) 2023 The INVRS-IO authors.
 """
 
 import json
-import requests
 import time
 from typing import Any, Dict, Optional
 
+import requests
 from totypes import json_utils
 
-from invrs_opt import base
 from invrs_opt.experimental import labels
-
+from invrs_opt.optimizers import base
 
 PyTree = Any
 StateToken = str
@@ -133,7 +132,11 @@ def optimizer_client(
         response = json.loads(get_response.text)
         return json_utils.pytree_from_json(response[labels.PARAMS])
 
-    return base.Optimizer(init=init_fn, update=update_fn, params=params_fn)
+    return base.Optimizer(
+        init=init_fn,
+        update=update_fn,  # type: ignore[arg-type]
+        params=params_fn,
+    )
 
 
 # -----------------------------------------------------------------------------

invrs_opt/{base.py → optimizers/base.py}

@@ -4,8 +4,13 @@ Copyright (c) 2023 The INVRS-IO authors.
 """
 
 import dataclasses
+import inspect
 from typing import Any, Protocol
 
+import jax.numpy as jnp
+import optax  # type: ignore[import-untyped]
+from totypes import json_utils
+
 PyTree = Any
 
 
@@ -30,7 +35,7 @@ class UpdateFn(Protocol):
         self,
         *,
         grad: PyTree,
-        value: float,
+        value: jnp.ndarray,
         params: PyTree,
         state: PyTree,
     ) -> PyTree:
@@ -44,3 +49,13 @@ class Optimizer:
     init: InitFn
     params: ParamsFn
     update: UpdateFn
+
+
+# Register all optax state types for serialization.
+optax_types = {}
+for name, obj in inspect.getmembers(optax):
+    if name.endswith("State") and isinstance(obj, type):
+        optax_types[obj] = True
+
+for obj in optax_types.keys():
+    json_utils.register_custom_type(obj)