Convert this array to numpy array Returns ------- np_arr : numpy.ndarray The corresponding numpy array.
(self)
| 167 | return str(self.numpy()) |
| 168 | |
| 169 | def numpy(self): |
| 170 | """Convert this array to numpy array |
| 171 | |
| 172 | Returns |
| 173 | ------- |
| 174 | np_arr : numpy.ndarray |
| 175 | The corresponding numpy array. |
| 176 | """ |
| 177 | t = tvm_ffi.dtype(self.dtype) |
| 178 | shape, dtype = self.shape, self.dtype |
| 179 | old_dtype = dtype |
| 180 | if t.lanes > 1: |
| 181 | shape = shape + (t.lanes,) |
| 182 | t = t.with_lanes(1) |
| 183 | dtype = str(t) |
| 184 | if dtype == "int4": |
| 185 | dtype = "int8" |
| 186 | if dtype in [ |
| 187 | "bfloat16", |
| 188 | "float8_e3m4", |
| 189 | "float8_e4m3", |
| 190 | "float8_e4m3b11fnuz", |
| 191 | "float8_e4m3fn", |
| 192 | "float8_e4m3fnuz", |
| 193 | "float8_e5m2", |
| 194 | "float8_e5m2fnuz", |
| 195 | "float8_e8m0fnu", |
| 196 | "float6_e2m3fn", |
| 197 | "float6_e3m2fn", |
| 198 | "float4_e2m1fn", |
| 199 | ]: |
| 200 | if ml_dtypes is None: |
| 201 | raise RuntimeError( |
| 202 | f"ml_dtypes is not installed, cannot convert {dtype} array to numpy." |
| 203 | ) |
| 204 | try: |
| 205 | dtype = getattr(ml_dtypes, dtype) |
| 206 | except AttributeError: |
| 207 | raise RuntimeError(f"ml_dtypes has no attribute '{dtype}', cannot convert array.") |
| 208 | np_arr = np.empty(shape, dtype=dtype) |
| 209 | assert np_arr.flags["C_CONTIGUOUS"] |
| 210 | data = np_arr.ctypes.data_as(ctypes.c_void_p) |
| 211 | # TODO(kathy): revisit and get a mirrored function of ffi::GetDataSize |
| 212 | # in Python to replace line below |
| 213 | nbytes = np_arr.size if dtype == "bool" else (np_arr.size * old_dtype.bits + 7) // 8 |
| 214 | _ffi_api.TVMTensorCopyToBytes(self, data, nbytes) |
| 215 | |
| 216 | if old_dtype == "int4" or old_dtype.startswith("float4_e2m1fn"): |
| 217 | length = np_arr.size |
| 218 | np_arr = np_arr.view("int8") |
| 219 | np_arr_ret = np.empty((length,), dtype="int8") |
| 220 | np_arr = np_arr.reshape((length,)) |
| 221 | odd_index = np.bitwise_and(np_arr, 0x0F) |
| 222 | even_index = np.bitwise_and(np_arr >> 4, 0x0F) |
| 223 | np_arr_ret[1::2] = odd_index[0 : length // 2] |
| 224 | np_arr_ret[0::2] = even_index[0 : (length + 1) // 2] |
| 225 | return np_arr_ret.reshape(shape).view(dtype) |
| 226 |