Remove dependency of iree.runtime to iree.runtime.distributed

nod-ai · Aug 15, 2023 · c0e9d82 · c0e9d82
1 parent 3f1e5e7
commit c0e9d82
Show file tree

Hide file tree

Showing 4 changed files with 1 addition and 4 deletions.
diff --git a/runtime/bindings/python/iree/runtime/__init__.py b/runtime/bindings/python/iree/runtime/__init__.py
@@ -57,5 +57,4 @@
 from .function import *
 from .tracing import *
 
-from . import distributed
 from . import flags
diff --git a/runtime/bindings/python/iree/runtime/distributed/distributed.py b/runtime/bindings/python/iree/runtime/distributed/distributed.py
@@ -4,7 +4,6 @@
 # See https://llvm.org/LICENSE.txt for license information.
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
-import iree.compiler
 import sys
 import iree.runtime
 from iree.runtime.array_interop import DeviceArray

diff --git a/runtime/bindings/python/iree/runtime/distributed/run_rank.py b/runtime/bindings/python/iree/runtime/distributed/run_rank.py
@@ -4,7 +4,6 @@
 # See https://llvm.org/LICENSE.txt for license information.
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
-import iree.compiler
 import argparse
 import iree.runtime
 from iree.runtime.array_interop import DeviceArray

diff --git a/runtime/bindings/python/iree/runtime/distributed/sharding_pass_validation.py b/runtime/bindings/python/iree/runtime/distributed/sharding_pass_validation.py
@@ -1,7 +1,7 @@
 import iree.compiler
 import iree.runtime
 import os
-from iree.runtime.distributed import run_ranks
+from .distributed import run_ranks
 import subprocess
 from pathlib import Path
 from jax._src.lib import xla_client