Add AWS Trainium device

Author: ndgrigorian

dpctl/tensor/_dlpack.pxd

@@ -41,6 +41,7 @@ cdef extern from "dlpack/dlpack.h" nogil:
     int device_WebGPU "kDLWebGPU"
     int device_Hexagon "kDLHexagon"
     int device_MAIA "kDLMAIA"
+    int device_Trn "kDLTrn"
 
 cpdef object to_dlpack_capsule(usm_ndarray array) except +
 cpdef object to_dlpack_versioned_capsule(

dpctl/tensor/_dlpack.pyx

@@ -76,6 +76,7 @@ cdef extern from "dlpack/dlpack.h" nogil:
         kDLWebGPU
         kDLHexagon
         kDLMAIA
+        kDLTrn
 
     ctypedef struct DLDevice:
         DLDeviceType device_type

dpctl/tensor/_usmarray.pyx

@@ -86,6 +86,8 @@ class DLDeviceType(IntEnum):
             Qualcomm Hexagon DSP
         ``kDLMAIA``:
             Microsoft MAIA device
+        ``kDLTrn``:
+            AWS Trainium device
     """
     kDLCPU = c_dlpack.device_CPU
     kDLCUDA = c_dlpack.device_CUDA
@@ -101,6 +103,7 @@ class DLDeviceType(IntEnum):
     kDLWebGPU = c_dlpack.device_WebGPU
     kDLHexagon = c_dlpack.device_Hexagon
     kDLMAIA = c_dlpack.device_MAIA
+    kDLTrn = c_dlpack.device_Trn
 
 
 cdef class InternalUSMArrayError(Exception):