Unverified Commit 12c40963 authored by Matthew Douglas's avatar Matthew Douglas Committed by GitHub
Browse files

Add autoloading for backend packages (#1593)


Co-authored-by: default avatarTitus von Koeller <9048635+Titus-von-Koeller@users.noreply.github.com>
parent 8f88cef1
...@@ -4,6 +4,8 @@ ...@@ -4,6 +4,8 @@
# LICENSE file in the root directory of this source tree. # LICENSE file in the root directory of this source tree.
import sys
import torch import torch
from . import _ops, research, utils from . import _ops, research, utils
...@@ -32,6 +34,30 @@ supported_torch_devices = { ...@@ -32,6 +34,30 @@ supported_torch_devices = {
if torch.cuda.is_available(): if torch.cuda.is_available():
from .backends.cuda import ops as cuda_ops from .backends.cuda import ops as cuda_ops
def _import_backends():
"""
Discover and autoload all available backends installed as separate packages.
Packages with an entrypoint for "bitsandbytes.backends" will be loaded.
Inspired by PyTorch implementation: https://pytorch.org/tutorials/prototype/python_extension_autoload.html
"""
from importlib.metadata import entry_points
if sys.version_info < (3, 10):
extensions = entry_points().get("bitsandbytes.backends", [])
else:
extensions = entry_points(group="bitsandbytes.backends")
for ext in extensions:
try:
entry = ext.load()
entry()
except Exception as e:
raise RuntimeError(f"bitsandbytes: failed to load backend {ext.name}: {e}") from e
_import_backends()
__pdoc__ = { __pdoc__ = {
"libbitsandbytes": False, "libbitsandbytes": False,
"optim.optimizer.Optimizer8bit": False, "optim.optimizer.Optimizer8bit": False,
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment