(A, B, out, transposed_A, transposed_B, expected_type=torch.int8)
| 1214 | |
| 1215 | @deprecated("This function is deprecated and will be removed in a future release.", category=FutureWarning) |
| 1216 | def check_matmul(A, B, out, transposed_A, transposed_B, expected_type=torch.int8): |
| 1217 | if not torch.cuda.is_initialized(): |
| 1218 | torch.cuda.init() |
| 1219 | if A.dtype != expected_type or B.dtype != expected_type: |
| 1220 | raise TypeError(f"Expected torch.int8 input tensors A and B, but got {A.dtype} and {B.dtype}") |
| 1221 | |
| 1222 | sA = A.shape |
| 1223 | sB = B.shape |
| 1224 | tA = transposed_A |
| 1225 | tB = transposed_B |
| 1226 | |
| 1227 | correct = True |
| 1228 | |
| 1229 | if len(sA) == 2 and len(sB) == 2: |
| 1230 | if not tA and not tB and A.shape[1] != B.shape[0]: |
| 1231 | correct = False |
| 1232 | elif tA and not tB and A.shape[0] != B.shape[0]: |
| 1233 | correct = False |
| 1234 | elif tA and tB and A.shape[0] != B.shape[1]: |
| 1235 | correct = False |
| 1236 | elif not tA and tB and A.shape[1] != B.shape[1]: |
| 1237 | correct = False |
| 1238 | elif len(sA) == 3 and len(sB) == 2: |
| 1239 | if not tA and not tB and A.shape[2] != B.shape[0]: |
| 1240 | correct = False |
| 1241 | elif tA and not tB and A.shape[1] != B.shape[0]: |
| 1242 | correct = False |
| 1243 | elif tA and tB and A.shape[1] != B.shape[1]: |
| 1244 | correct = False |
| 1245 | elif not tA and tB and A.shape[2] != B.shape[1]: |
| 1246 | correct = False |
| 1247 | elif len(sA) == 3 and len(sB) == 3: |
| 1248 | if not tA and not tB and A.shape[2] != B.shape[1]: |
| 1249 | correct = False |
| 1250 | elif tA and not tB and A.shape[1] != B.shape[1]: |
| 1251 | correct = False |
| 1252 | elif tA and tB and A.shape[1] != B.shape[2]: |
| 1253 | correct = False |
| 1254 | elif not tA and tB and A.shape[2] != B.shape[2]: |
| 1255 | correct = False |
| 1256 | |
| 1257 | if out is not None: |
| 1258 | sout = out.shape |
| 1259 | # special case common in backprop |
| 1260 | if not correct and len(sA) == 3 and len(sB) == 3: |
| 1261 | if sout[0] == sA[2] and sout[1] == sB[2] and sA[0] == sB[0] and sA[1] == sB[1]: |
| 1262 | correct = True |
| 1263 | else: |
| 1264 | if len(sA) == 2 and len(sB) == 2: |
| 1265 | if not tA and not tB: |
| 1266 | sout = (sA[0], sB[1]) |
| 1267 | elif tA and tB: |
| 1268 | sout = (sA[1], sB[0]) |
| 1269 | elif tA and not tB: |
| 1270 | sout = (sA[1], sB[1]) |
| 1271 | elif not tA and tB: |
| 1272 | sout = (sA[0], sB[0]) |
| 1273 | elif len(sA) == 3 and len(sB) == 2: |
no test coverage detected