|
|
import importlib |
|
|
|
|
|
|
|
|
class LazyModule: |
|
|
def __init__(self, name, pip_name=None, import_error_msg=None): |
|
|
self.name = name |
|
|
self.pip_name = pip_name or name |
|
|
self.import_error_msg = import_error_msg or ( |
|
|
f"This requires the {self.name} module. " |
|
|
f"You can install it via `pip install {self.pip_name}`" |
|
|
) |
|
|
self.module = None |
|
|
self._available = None |
|
|
|
|
|
@property |
|
|
def available(self): |
|
|
if self._available is None: |
|
|
try: |
|
|
self.initialize() |
|
|
self._available = True |
|
|
except ImportError: |
|
|
self._available = False |
|
|
return self._available |
|
|
|
|
|
def initialize(self): |
|
|
try: |
|
|
self.module = importlib.import_module(self.name) |
|
|
except ImportError: |
|
|
raise ImportError(self.import_error_msg) |
|
|
|
|
|
def __getattr__(self, name): |
|
|
if name == "_api_export_path": |
|
|
raise AttributeError |
|
|
if self.module is None: |
|
|
self.initialize() |
|
|
return getattr(self.module, name) |
|
|
|
|
|
def __repr__(self): |
|
|
return f"LazyModule({self.name})" |
|
|
|
|
|
|
|
|
tensorflow = LazyModule("tensorflow") |
|
|
gfile = LazyModule("tensorflow.io.gfile", pip_name="tensorflow") |
|
|
tensorflow_io = LazyModule("tensorflow_io") |
|
|
scipy = LazyModule("scipy") |
|
|
jax = LazyModule("jax") |
|
|
torch_xla = LazyModule( |
|
|
"torch_xla", |
|
|
import_error_msg=( |
|
|
"This requires the torch_xla module. You can install it via " |
|
|
"`pip install torch-xla`. Additionally, you may need to update " |
|
|
"LD_LIBRARY_PATH if necessary. Torch XLA builds a shared library, " |
|
|
"_XLAC.so, which needs to link to the version of Python it was built " |
|
|
"with. Use the following command to update LD_LIBRARY_PATH: " |
|
|
"`export LD_LIBRARY_PATH=<path to Python>/lib:$LD_LIBRARY_PATH`" |
|
|
), |
|
|
) |
|
|
optree = LazyModule("optree") |
|
|
dmtree = LazyModule("tree") |
|
|
tf2onnx = LazyModule("tf2onnx") |
|
|
|
