Merge branch 'google:main' into main

muralidhar-challa · web-flow · commit 269d06ceff6a · 2026-02-17T09:18:10.000+05:30
diff --git a/docs/examples/jaxonnxruntime-python-api.py b/docs/examples/jaxonnxruntime-python-api.py
@@ -26,11 +26,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """Jaxonnxruntime python API example."""
+
 from jaxonnxruntime import backend as jax_backend
 import numpy as np
-import torch
-
 import onnx
+import torch
 
 MODEL_FILE = ".model.onnx"
 
diff --git a/jaxonnxruntime/backend.py b/jaxonnxruntime/backend.py
@@ -31,7 +31,6 @@
 
 import jax
 from jaxonnxruntime.core import call_onnx
-
 import onnx
 
 
diff --git a/jaxonnxruntime/core/call_onnx.py b/jaxonnxruntime/core/call_onnx.py
@@ -24,7 +24,6 @@
 from jaxonnxruntime.core import onnx_node
 from jaxonnxruntime.core import onnx_utils
 import jaxtyping
-
 import onnx
 from onnx import helper as onnx_helper
 
diff --git a/jaxonnxruntime/core/call_onnx_test.py b/jaxonnxruntime/core/call_onnx_test.py
@@ -18,7 +18,6 @@
 from jaxonnxruntime.core import config_class
 from jaxonnxruntime.onnx_ops import abs
 import numpy as np
-
 import onnx
 
 
diff --git a/jaxonnxruntime/core/onnx_graph.py b/jaxonnxruntime/core/onnx_graph.py
@@ -35,7 +35,6 @@
 from jaxonnxruntime.core import onnx_node
 from jaxonnxruntime.core import onnx_utils
 import numpy as np
-
 import onnx
 
 
diff --git a/jaxonnxruntime/core/onnx_graph_test.py b/jaxonnxruntime/core/onnx_graph_test.py
@@ -14,7 +14,6 @@
 
 from absl.testing import absltest
 from jaxonnxruntime.core import onnx_graph
-
 import onnx
 
 OnnxGraph = onnx_graph.OnnxGraph
diff --git a/jaxonnxruntime/core/onnx_node_test.py b/jaxonnxruntime/core/onnx_node_test.py
@@ -14,7 +14,6 @@
 
 from absl.testing import absltest
 from jaxonnxruntime.core import onnx_node
-
 import onnx
 
 
diff --git a/jaxonnxruntime/core/onnx_utils.py b/jaxonnxruntime/core/onnx_utils.py
@@ -23,7 +23,6 @@
 from jax import numpy as jnp
 from jaxonnxruntime.core import config_class
 import numpy as np
-
 import onnx
 from onnx import helper as onnx_helper
 from onnx import numpy_helper
diff --git a/jaxonnxruntime/core/onnx_utils_test.py b/jaxonnxruntime/core/onnx_utils_test.py
@@ -14,7 +14,6 @@
 
 from absl.testing import absltest
 from jaxonnxruntime.core import onnx_utils
-
 import onnx
 
 
diff --git a/jaxonnxruntime/experimental/call_torch/_call_torch.py b/jaxonnxruntime/experimental/call_torch/_call_torch.py
@@ -21,9 +21,8 @@
 from absl import logging
 import jax
 from jaxonnxruntime.core import call_onnx
-import torch
-
 import onnx
+import torch
 
 
 class TorchONNXExportError(Exception):
diff --git a/jaxonnxruntime/experimental/call_torch/call_torch_test_util.py b/jaxonnxruntime/experimental/call_torch/call_torch_test_util.py
@@ -24,9 +24,8 @@
 from jaxonnxruntime.core import onnx_utils
 from jaxonnxruntime.experimental import call_torch
 import numpy as np
-import torch
-
 import onnx
+import torch
 
 
 def _maybe_upcast(x: jax.Array):
diff --git a/jaxonnxruntime/onnx_ops/if_op.py b/jaxonnxruntime/onnx_ops/if_op.py
@@ -39,7 +39,6 @@
 from jaxonnxruntime.core import handler
 from jaxonnxruntime.core import onnx_node
 from jaxonnxruntime.core import onnx_utils
-
 import onnx
 
 config = config_class.config
diff --git a/jaxonnxruntime/onnx_ops/onehot_test.py b/jaxonnxruntime/onnx_ops/onehot_test.py
@@ -16,7 +16,6 @@
 from jaxonnxruntime.core import onnx_node
 from jaxonnxruntime.onnx_ops import onehot
 import numpy as np
-
 import onnx
 
 NodeProto = onnx.NodeProto
diff --git a/jaxonnxruntime/onnx_ops/quantizelinear.py b/jaxonnxruntime/onnx_ops/quantizelinear.py
@@ -25,7 +25,6 @@
 from jaxonnxruntime.core import onnx_node
 from jaxonnxruntime.core import onnx_utils
 from jaxonnxruntime.onnx_ops import onnx_ops_utils
-
 import onnx
 
 
diff --git a/jaxonnxruntime/runner.py b/jaxonnxruntime/runner.py
@@ -31,7 +31,6 @@
 import jax
 from jaxonnxruntime import backend as jort_backend
 import numpy as np
-
 import onnx
 from onnx import numpy_helper
 
@@ -68,7 +67,9 @@ class TestCase:
 
 
 DATA_DIR = os.path.join(
-    os.path.realpath(os.path.dirname(onnx.__file__)), 'backend/test/data'
+    # TODO(derricktli): Fix this hack.
+    os.path.realpath(os.path.dirname(numpy_helper.__file__)),
+    'backend/test/data',
 )
 
 
diff --git a/tests/onnx_models_test.py b/tests/onnx_models_test.py
@@ -27,7 +27,6 @@
 from jaxonnxruntime.core import onnx_utils
 from onnx import hub
 import numpy as np
-
 import onnx
 
 config = jaxonnxruntime.config
