diff --git a/.gitignore b/.gitignore
index 0022aeb42..031ed9eeb 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,6 +1,7 @@
 # Xmake cache
 .xmake/
 build/
+python/infinicore/lib/*.so
 
 # MacOS Cache
 .DS_Store
diff --git a/python/infinicore/device.py b/python/infinicore/device.py
index c54d5d71f..acb835969 100644
--- a/python/infinicore/device.py
+++ b/python/infinicore/device.py
@@ -1,4 +1,4 @@
-from . import _infinicore
+from infinicore.lib import _infinicore
 
 
 class device:
diff --git a/python/infinicore/dtype.py b/python/infinicore/dtype.py
index d85f97f87..ec06f5d7f 100644
--- a/python/infinicore/dtype.py
+++ b/python/infinicore/dtype.py
@@ -1,4 +1,4 @@
-from . import _infinicore
+from infinicore.lib import _infinicore
 
 
 class dtype:
diff --git a/python/infinicore/ops/matmul.py b/python/infinicore/ops/matmul.py
index b0c167823..166ccdaa8 100644
--- a/python/infinicore/ops/matmul.py
+++ b/python/infinicore/ops/matmul.py
@@ -1,7 +1,6 @@
+from infinicore.lib import _infinicore
 from infinicore.tensor import Tensor
 
-from .. import _infinicore
-
 
 def matmul(input, other, *, out=None):
     if out is None:
diff --git a/python/infinicore/ops/rearrange.py b/python/infinicore/ops/rearrange.py
index d718d3658..c53f62bc8 100644
--- a/python/infinicore/ops/rearrange.py
+++ b/python/infinicore/ops/rearrange.py
@@ -1,7 +1,6 @@
+from infinicore.lib import _infinicore
 from infinicore.tensor import Tensor
 
-from .. import _infinicore
-
 
 def rearrange(input, other, *, out=None):
     if out is None:
diff --git a/python/infinicore/tensor.py b/python/infinicore/tensor.py
index 5592993e2..5095768c5 100644
--- a/python/infinicore/tensor.py
+++ b/python/infinicore/tensor.py
@@ -1,7 +1,6 @@
 import infinicore.device
 import infinicore.dtype
-
-from . import _infinicore
+from infinicore.lib import _infinicore
 
 
 class Tensor:
diff --git a/setup.py b/setup.py
index d46d634c7..8a2ccc34f 100644
--- a/setup.py
+++ b/setup.py
@@ -1,8 +1,4 @@
-import glob
-import os
-import shutil
 import subprocess
-from pathlib import Path
 
 from setuptools import setup
 from setuptools.command.build import build
@@ -15,12 +11,5 @@ def run(self):
         subprocess.run(["xmake", "build", "-y", "_infinicore"])
         subprocess.run(["xmake", "install", "_infinicore"])
 
-        installation_dir = os.getenv("INFINI_ROOT", str(Path.home() / ".infini"))
-        lib_dir = os.path.join(installation_dir, "lib")
-        lib_path = glob.glob(os.path.join(lib_dir, "_infinicore.*"))[0]
-        package_dir = os.path.join(self.build_lib, "infinicore")
-        os.makedirs(package_dir, exist_ok=True)
-        shutil.move(lib_path, package_dir)
-
 
 setup(package_dir={"": "python"}, cmdclass={"build": Build})
diff --git a/xmake.lua b/xmake.lua
index df3e39faf..e517aeec2 100644
--- a/xmake.lua
+++ b/xmake.lua
@@ -347,7 +347,31 @@ target("_infinicore")
     add_files("src/infinicore/op/*/*.cc")
     add_files("src/infinicore/pybind11/**.cc")
 
-    set_installdir(os.getenv("INFINI_ROOT") or (os.getenv(is_host("windows") and "HOMEPATH" or "HOME") .. "/.infini"))
+    set_installdir("python/infinicore")
+target_end()
+
+option("editable")
+    set_default(false)
+    set_showmenu(true)
+    set_description("Install the `infinicore` Python package in editable mode")
+option_end()
+
+target("infinicore")
+    set_kind("phony")
+
+    set_default(false)
+
+    add_deps("_infinicore")
+
+    on_install(function (target)
+        local pip_install_args = {}
+
+        if has_config("editable") then
+            table.insert(pip_install_args, "--editable")
+        end
+
+        os.execv("python", table.join({"-m", "pip", "install"}, pip_install_args, {"."}))
+    end)
 target_end()
 
 -- Tests