imt-ring 1.6.24__tar.gz → 1.6.26__tar.gz

{imt_ring-1.6.24 → imt_ring-1.6.26}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: imt-ring
-Version: 1.6.24
+Version: 1.6.26
 Summary: RING: Recurrent Inertial Graph-based Estimator
 Author-email: Simon Bachhuber <simon.bachhuber@fau.de>
 Project-URL: Homepage, https://github.com/SimiPixel/ring

{imt_ring-1.6.24 → imt_ring-1.6.26}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "imt-ring"
-version = "1.6.24"
+version = "1.6.26"
 authors = [
   { name="Simon Bachhuber", email="simon.bachhuber@fau.de" },
 ]

{imt_ring-1.6.24 → imt_ring-1.6.26}/src/imt_ring.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: imt-ring
-Version: 1.6.24
+Version: 1.6.26
 Summary: RING: Recurrent Inertial Graph-based Estimator
 Author-email: Simon Bachhuber <simon.bachhuber@fau.de>
 Project-URL: Homepage, https://github.com/SimiPixel/ring

{imt_ring-1.6.24 → imt_ring-1.6.26}/src/ring/ml/ml_utils.py

@@ -243,5 +243,102 @@ def save_model_tf(jax_func, path: str, *input, validate: bool = True):
         )
 
 
+def to_onnx(
+    fn,
+    output_path,
+    *args: tuple[np.ndarray],
+    in_args_names: Optional[list[str]] = None,
+    out_args_names: Optional[list[str]] = None,
+    validate: bool = False,
+):
+    """
+    Converts a JAX function to ONNX format, with optional input/output renaming and validation.
+
+    Args:
+        fn (callable): The JAX function to be converted.
+        output_path (str): Path where the ONNX model will be saved.
+        *args (tuple[np.ndarray]): Input arguments for the JAX function.
+        in_args_names (Optional[list[str]]): Names for the ONNX model's input tensors. Defaults to None.
+        out_args_names (Optional[list[str]]): Names for the ONNX model's output tensors. Defaults to None.
+        validate (bool): Whether to validate the ONNX model against the JAX function's outputs. Defaults to False.
+
+    Raises:
+        AssertionError: If the number of provided names does not match the number of inputs/outputs.
+        AssertionError: If the ONNX model's outputs do not match the JAX function's outputs within tolerance.
+        ValueError: If any error occurs during ONNX conversion, saving, or validation.
+
+    Notes:
+        - The function uses `jax2tf` to convert the JAX function to TensorFlow format,
+          and `tf2onnx` for ONNX conversion.
+        - Input and output tensor names in the ONNX model can be renamed using `sor4onnx.rename`.
+        - Validation compares outputs of the JAX function and the ONNX model using ONNX Runtime.
+
+    Example:
+        ```
+        import jax.numpy as jnp
+
+        def my_fn(x, y):
+            return x + y, x * y
+
+        x = jnp.array([1, 2, 3])
+        y = jnp.array([4, 5, 6])
+
+        to_onnx(
+            my_fn,
+            "model.onnx",
+            x, y,
+            in_args_names=["input1", "input2"],
+            out_args_names=["sum", "product"],
+            validate=True,
+        )
+        ```
+    """  # noqa: E501
+    import jax.experimental.jax2tf as jax2tf
+    import tensorflow as tf
+    import tf2onnx
+
+    tf_fn = tf.function(jax2tf.convert(fn, enable_xla=False))
+    tf_args = [tf.TensorSpec(np.shape(x), np.result_type(x)) for x in args]
+    tf2onnx.convert.from_function(
+        tf_fn, input_signature=tf_args, output_path=output_path
+    )
+
+    if in_args_names is not None or out_args_names is not None:
+        import onnx
+        from sor4onnx import rename
+
+        model = onnx.load(output_path)
+
+        if in_args_names is not None:
+            old_names = [inp.name for inp in model.graph.input]
+            assert len(old_names) == len(in_args_names)
+            for old_name, new_name in zip(old_names, in_args_names):
+                model = rename([old_name, new_name], None, model, None, mode="inputs")
+
+        if out_args_names is not None:
+            old_names = [out.name for out in model.graph.output]
+            assert len(old_names) == len(out_args_names)
+            for old_name, new_name in zip(old_names, out_args_names):
+                model = rename([old_name, new_name], None, model, None, mode="outputs")
+
+        onnx.save(model, output_path)
+
+    if validate:
+        import onnxruntime as ort
+
+        output_jax = fn(*args)
+        session = ort.InferenceSession(output_path)
+        input_names = [inp.name for inp in session.get_inputs()]
+        output_onnx = session.run(
+            None, {name: np.array(arg) for name, arg in zip(input_names, args)}
+        )
+
+        for o1, o2 in zip(output_jax, output_onnx):
+            assert np.allclose(o1, o2, atol=1e-5, rtol=1e-5)
+
+        if out_args_names is not None:
+            assert [out.name for out in session.get_outputs()] == out_args_names
+
+
 def _unknown_link_names(N: int):
     return [f"link{i}" for i in range(N)]

{imt_ring-1.6.24 → imt_ring-1.6.26}/src/ring/ml/train.py

@@ -45,7 +45,7 @@ def _build_step_fn(
 
     @partial(jax.value_and_grad, has_aux=True)
     def loss_fn(params, state, X, y):
-        yhat, state = filter.apply(params=params, state=state, X=X)
+        yhat, state = filter.apply(params=params, state=state, X=X, y=y)
         # this vmap maps along batch-axis, not time-axis
         # time-axis is handled by `metric_fn`
         pipe = lambda q, qhat: jnp.mean(jax.vmap(metric_fn)(q, qhat))
@@ -261,7 +261,7 @@ def _build_eval_fn(
     """Build function that evaluates the filter performance."""
 
     def eval_fn(params, state, X, y):
-        yhat, _ = filter.apply(params=params, state=state, X=X)
+        yhat, _ = filter.apply(params=params, state=state, X=X, y=y)
 
         y = _arr_to_dict(y, link_names)
         yhat = _arr_to_dict(yhat, link_names)