llvm
/
torch-mlir
mirror of https://github.com/llvm/torch-mlir
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Yapf Format `refbacked.py`.

pull/1727/head
Prashant Kumar 2022-12-14 15:19:26 +00:00
parent 564403e3a1
commit 8ba77ae2a5
1 changed files with 7 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
python/torch_mlir_e2e_test/linalg_on_tensors_backends/refbackend.py
View File

@ -21,7 +21,10 @@ __all__ = [
def assert_arg_type_is_supported(ty):
SUPPORTED = [np.float16, np.float32, np.float64, np.uint8, np.int8, np.int32, np.int64, np.bool_]
SUPPORTED = [
np.float16, np.float32, np.float64, np.uint8, np.int8, np.int32,
np.int64, np.bool_
]
assert ty in SUPPORTED, f"Only numpy arrays with dtypes in {SUPPORTED} are supported, but got {ty}"
@ -86,8 +89,9 @@ class RefBackendInvoker:
def consume_return_funcs(*args):
self.result = tuple([
arg if type in elemental_type_to_ctype else
unranked_memref_to_numpy(arg, memref_type_to_np_dtype[type])
arg if type in elemental_type_to_ctype
else unranked_memref_to_numpy(
arg, memref_type_to_np_dtype[type])
for arg, type in zip(args, ret_types)
])
if len(self.result) == 1: