Add TPU support (#629)

* tpu support 1

* change package name

* run format

* add torch_xla dependency

* run poetry lock --no-update

* Delete libs/infinity_emb/poetry.lock

* Update pyproject.toml

* Create poetry.lock

---------

Co-authored-by: Michael Feil <63565275+michaelfeil@users.noreply.github.com>

Author: wirthual

libs/infinity_emb/infinity_emb/_optional_imports.py

@@ -78,3 +78,4 @@ def _raise_error(self) -> None:
 CHECK_TRANSFORMERS = OptionalImports("transformers", "torch")
 CHECK_TYPER = OptionalImports("typer", "server")
 CHECK_UVICORN = OptionalImports("uvicorn", "server")
+CHECK_XLA = OptionalImports("torch_xla", "torch_xla")

libs/infinity_emb/infinity_emb/inference/loading_strategy.py

@@ -1,4 +1,4 @@
-from infinity_emb._optional_imports import CHECK_TORCH, CHECK_TRANSFORMERS
+from infinity_emb._optional_imports import CHECK_TORCH, CHECK_TRANSFORMERS, CHECK_XLA
 from infinity_emb.args import EngineArgs
 from infinity_emb.primitives import InferenceEngine, Device, Dtype, DeviceID, LoadingStrategy
 
@@ -7,6 +7,10 @@
     import torch
 if CHECK_TRANSFORMERS.is_available:
     from transformers import is_torch_npu_available  # type: ignore
+    from transformers.utils.import_utils import is_torch_xla_available # type: ignore
+
+if CHECK_XLA.is_available:
+    import torch_xla # type: ignore
 
 
 def _validate_availale_device_ids(
@@ -35,6 +39,8 @@ def get_loading_strategy_torch(args: EngineArgs) -> LoadingStrategy:
             autodevice = "npu"
         elif torch.backends.mps.is_available():
             autodevice = "mps"
+        elif is_torch_xla_available():
+            autodevice = "xla"
         else:
             autodevice = "cpu"
     else:
@@ -58,6 +64,10 @@ def get_loading_strategy_torch(args: EngineArgs) -> LoadingStrategy:
     elif autodevice == "cpu":
         # spawn multiple processes on CPU. This is useful for debugging, but not for performance.
         autodevice_string = ["cpu"] * max(len(args.device_id), 1)
+    elif autodevice == "xla":
+        autodevice_string = _validate_availale_device_ids(
+            "xla", list(range(torch_xla.device_count())), args.device_id
+        )
     else:
         raise ValueError(f"Unknown device {autodevice}")
 

libs/infinity_emb/infinity_emb/primitives.py

@@ -109,6 +109,7 @@ class Device(EnumType):
     cuda = "cuda"
     mps = "mps"
     tensorrt = "tensorrt"
+    xla = "xla"
     auto = "auto"
 
     @staticmethod