Checks if we can use BLAS (np.tensordot) call and its beneficial to do so. Parameters ---------- inputs : list of str Specifies the subscripts for summation. result : str Resulting summation. idx_removed : set Indices that are removed in the summatio
(inputs, result, idx_removed)
| 411 | |
| 412 | |
| 413 | def _can_dot(inputs, result, idx_removed): |
| 414 | """ |
| 415 | Checks if we can use BLAS (np.tensordot) call and its beneficial to do so. |
| 416 | |
| 417 | Parameters |
| 418 | ---------- |
| 419 | inputs : list of str |
| 420 | Specifies the subscripts for summation. |
| 421 | result : str |
| 422 | Resulting summation. |
| 423 | idx_removed : set |
| 424 | Indices that are removed in the summation |
| 425 | |
| 426 | |
| 427 | Returns |
| 428 | ------- |
| 429 | type : bool |
| 430 | Returns true if BLAS should and can be used, else False |
| 431 | |
| 432 | Notes |
| 433 | ----- |
| 434 | If the operations is BLAS level 1 or 2 and is not already aligned |
| 435 | we default back to einsum as the memory movement to copy is more |
| 436 | costly than the operation itself. |
| 437 | |
| 438 | |
| 439 | Examples |
| 440 | -------- |
| 441 | |
| 442 | # Standard GEMM operation |
| 443 | >>> _can_dot(['ij', 'jk'], 'ik', set('j')) |
| 444 | True |
| 445 | |
| 446 | # Can use the standard BLAS, but requires odd data movement |
| 447 | >>> _can_dot(['ijj', 'jk'], 'ik', set('j')) |
| 448 | False |
| 449 | |
| 450 | # DDOT where the memory is not aligned |
| 451 | >>> _can_dot(['ijk', 'ikj'], '', set('ijk')) |
| 452 | False |
| 453 | |
| 454 | """ |
| 455 | |
| 456 | # All `dot` calls remove indices |
| 457 | if len(idx_removed) == 0: |
| 458 | return False |
| 459 | |
| 460 | # BLAS can only handle two operands |
| 461 | if len(inputs) != 2: |
| 462 | return False |
| 463 | |
| 464 | input_left, input_right = inputs |
| 465 | |
| 466 | for c in set(input_left + input_right): |
| 467 | # can't deal with repeated indices on same input or more than 2 total |
| 468 | nl, nr = input_left.count(c), input_right.count(c) |
| 469 | if (nl > 1) or (nr > 1) or (nl + nr > 2): |
| 470 | return False |