This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm-ffi.git
The following commit(s) were added to refs/heads/main by this push:
new ffa2dbf feat: add trainium device (#76)
ffa2dbf is described below
commit ffa2dbf8bc18edb3114f18f619da08c4e3289de6
Author: Bohan Hou <[email protected]>
AuthorDate: Wed Oct 1 08:52:27 2025 -0400
feat: add trainium device (#76)
---
python/tvm_ffi/core.pyi | 1 +
python/tvm_ffi/cython/device.pxi | 3 +++
2 files changed, 4 insertions(+)
diff --git a/python/tvm_ffi/core.pyi b/python/tvm_ffi/core.pyi
index ac04114..0c90ebf 100644
--- a/python/tvm_ffi/core.pyi
+++ b/python/tvm_ffi/core.pyi
@@ -143,6 +143,7 @@ class DLDeviceType(IntEnum):
kDLOneAPI = 14
kDLWebGPU = 15
kDLHexagon = 16
+ kDLTrn = 17
class Device:
"""Device represents a device in the ffi system."""
diff --git a/python/tvm_ffi/cython/device.pxi b/python/tvm_ffi/cython/device.pxi
index 84c047f..f25ee94 100644
--- a/python/tvm_ffi/cython/device.pxi
+++ b/python/tvm_ffi/cython/device.pxi
@@ -49,6 +49,7 @@ class DLDeviceType(IntEnum):
kDLOneAPI = 14
kDLWebGPU = 15
kDLHexagon = 16
+ kDLTrn = 17
cdef class Device:
@@ -90,6 +91,7 @@ cdef class Device:
DLDeviceType.kDLOneAPI: "oneapi",
DLDeviceType.kDLWebGPU: "webgpu",
DLDeviceType.kDLHexagon: "hexagon",
+ DLDeviceType.kDLTrn: "trn",
}
_DEVICE_NAME_TO_TYPE = {
@@ -108,6 +110,7 @@ cdef class Device:
"ext_dev": DLDeviceType.kDLExtDev,
"hexagon": DLDeviceType.kDLHexagon,
"webgpu": DLDeviceType.kDLWebGPU,
+ "trn": DLDeviceType.kDLTrn,
}
def __init__(self, device_type, index = None):
