imt-ring 1.6.24__py3-none-any.whl → 1.6.25__py3-none-any.whl

{imt_ring-1.6.24.dist-info → imt_ring-1.6.25.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: imt-ring
-Version: 1.6.24
+Version: 1.6.25
 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.dist-info → imt_ring-1.6.25.dist-info}/RECORD

@@ -54,7 +54,7 @@ ring/io/xml/to_xml.py,sha256=fohb-jWMf2cxVdT5dmknsGyrNMseICSbKEz_urbaWbQ,3407
 ring/ml/__init__.py,sha256=nbh48gaswWeY4S4vT1sply_3ROj2DQ7agjoLR4Ho3T8,1517
 ring/ml/base.py,sha256=lfwEZLBDglOSRWChUHoH1kezefhttPV9TMEpNIqsMNw,9972
 ring/ml/callbacks.py,sha256=W19QF6_uvaNCjs8ObsjNXD7mv9gFgJBixdRSbB_BynE,13301
-ring/ml/ml_utils.py,sha256=xqy9BnLy8IKVqkFS9mlZsGJXSbThI9zZxZ5rhl8LSI8,7144
+ring/ml/ml_utils.py,sha256=uAQ6qXFT2UxILwbKzFNPxaBeY4X56l9pixdv91MKQis,9072
 ring/ml/optimizer.py,sha256=TZF0_LmnewzmGVso-zIQJtpWguUW0fW3HeRpIdG_qoI,4763
 ring/ml/ringnet.py,sha256=mef7jyN2QcApJmQGH3HYZyTV-00q8YpsYOKhW0-ku1k,8973
 ring/ml/rnno_v1.py,sha256=2qE08OIvTJ5PvSxKpYGzGSrvEImWrdAT_qslZ7jP5tA,1372
@@ -86,7 +86,7 @@ ring/utils/randomize_sys.py,sha256=G_vBIo0OwQkXL2u0djwbaoaeb02C4LQCTNNloOYIU2M,3
 ring/utils/utils.py,sha256=tJaWXLGOTwkxJQj2l23dX97wO3aZYhM2qd7eNuMRs84,6907
 ring/utils/register_gym_envs/__init__.py,sha256=PtPIRBQJ16339xZ9G9VpvqrvcGbQ_Pk_SUz4tQPa9nQ,94
 ring/utils/register_gym_envs/saddle.py,sha256=tA5CyW_akSXyDm0xJ83CtOrUMVElH0f9vZtEDDJQalI,4422
-imt_ring-1.6.24.dist-info/METADATA,sha256=vaXarRf1r5xZeGK-av_regQ2LgaCTnb0Th43bDLXgN8,4089
-imt_ring-1.6.24.dist-info/WHEEL,sha256=PZUExdf71Ui_so67QXpySuHtCi3-J3wvF4ORK6k_S8U,91
-imt_ring-1.6.24.dist-info/top_level.txt,sha256=EiT790-lAyi8iwTzJArH3f2k77rwhDn00q-4PlmvDQo,5
-imt_ring-1.6.24.dist-info/RECORD,,
+imt_ring-1.6.25.dist-info/METADATA,sha256=8-77JWmLIy6E3nJVd2VqfxwoKHt9b26ruipjMKR2K8I,4089
+imt_ring-1.6.25.dist-info/WHEEL,sha256=PZUExdf71Ui_so67QXpySuHtCi3-J3wvF4ORK6k_S8U,91
+imt_ring-1.6.25.dist-info/top_level.txt,sha256=EiT790-lAyi8iwTzJArH3f2k77rwhDn00q-4PlmvDQo,5
+imt_ring-1.6.25.dist-info/RECORD,,

ring/ml/ml_utils.py

@@ -243,5 +243,60 @@ 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,
+):
+    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)]