MAINT: address linting issues

ComPWA · Apr 5, 2024 · d6e2cae · d6e2cae
1 parent b7a156b
commit d6e2cae
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 3 deletions.
diff --git a/src/jax_benchmark/dot_product.py b/src/jax_benchmark/dot_product.py
@@ -42,7 +42,9 @@ def main() -> int:
     args = parser.parse_args()
     mute_warnings()
     benchmarks = run_benchmark(
-        shape=tuple(int(i) for i in args.shape.split("x")),
+        shape=tuple(  # pyright:ignore[reportArgumentType]
+            int(i) for i in args.shape.split("x")
+        ),
         repeat=args.repeat,
         number=args.number,
     )

diff --git a/src/jax_benchmark/io.py b/src/jax_benchmark/io.py
@@ -12,9 +12,9 @@
 
 
 def enable_x64() -> None:
-    from jax.config import config  # pyright: ignore[reportPrivateImportUsage]
+    import jax
 
-    config.update("jax_enable_x64", True)
+    jax.config.update("jax_enable_x64", True)
 
 
 def get_machine_info() -> dict: