[FIXED] Matrix multiplication with a transposed NumPy array using Numba JIT does not work

Issue

Environment

• OS: Windows 10
• Python version: 3.10
• Numba version: 0.57.0
• NumPy version: 1.24.3

Example

import numpy as np
from numba import njit

@njit
def matmul_transposed(a: np.ndarray, b: np.ndarray) -> np.ndarray:
    # return a @ b.T  # also tried this with a similar result, np.matmul seems to be unsupported by Numba
    return a.dot(b.transpose())

matmul_transposed(np.array([[1.0, 1.0]]), np.array([[1.0, 1.0]]))

Error

The above example raises an error

- Resolution failure for literal arguments:
No implementation of function Function(<function array_dot at 0x...>) found for signature:
 >>> array_dot(array(float64, 2d, C), array(float64, 2d, F))
There are 2 candidate implementations:
  - Of which 2 did not match due to:
  Overload in function 'array_dot': File: numba\np\arrayobj.py: Line 5929.
    With argument(s): '(array(float64, 2d, C), array(float64, 2d, F))':
   Rejected as the implementation raised a specific error:
     TypingError: Failed in nopython mode pipeline (step: nopython frontend)
   No implementation of function Function(<function dot at 0x...>) found for signature:
    >>> dot(array(float64, 2d, C), array(float64, 2d, F))
   There are 4 candidate implementations:
         - Of which 2 did not match due to:
         Overload in function 'dot_2': File: numba\np\linalg.py: Line 525.
           With argument(s): '(array(float64, 2d, C), array(float64, 2d, F))':
          Rejected as the implementation raised a specific error:
            LoweringError: Failed in nopython mode pipeline (step: native lowering)
          scipy 0.16+ is required for linear algebra
          
          File "[...]\numba\np\linalg.py", line 582:
                      def _dot2_codegen(context, builder, sig, args):
                          <source elided>
                  return lambda left, right: _impl(left, right)
                  ^
          
          During: lowering "$8call_function.3 = call $2load_deref.0(left, right, func=$2load_deref.0, args=[Var(left, linalg.py:582), Var(right, linalg.py:582)], kws=(), vararg=None, varkwarg=None, target=None)" at [...]\numba\np\linalg.py (582)
     raised from [...]\numba\core\errors.py:837
         - Of which 2 did not match due to:
         Overload in function 'dot_3': File: numba\np\linalg.py: Line 784.
           With argument(s): '(array(float64, 2d, C), array(float64, 2d, F))':
          Rejected as the implementation raised a specific error:
            TypingError: missing a required argument: 'out'
     raised from [...]\numba\core\typing\templates.py:784
   
   During: resolving callee type: Function(<function dot at 0x...>)
   During: typing of call at [...]\numba\np\arrayobj.py (5932)
   
   
   File "[...]\numba\np\arrayobj.py", line 5932:
       def dot_impl(arr, other):
           return np.dot(arr, other)
           ^
  raised from [...]\numba\core\typeinfer.py:1086
- Resolution failure for non-literal arguments:
None
During: resolving callee type: BoundFunction((<class 'numba.core.types.npytypes.Array'>, 'dot') for array(float64, 2d, C))
During: typing of call at [...]\example.py (7)
File "scratch_2.py", line 7:
def matmul_transposed(a: np.ndarray, b: np.ndarray) -> np.ndarray:
    <source elided>
    """Return a @ b.T"""
    return a.dot(b.transpose())
    ^

Interpretation

From the error message I concluded that Numba seems to transpose the array by changing its layout style from C to Fotran which is a cheap operation as it does not have to change the layout physically but than it seems to not know how to multiply the C-style array and the Fotrtran style array.

Question

Is there any way to multiply these matrices? Preferably without copying the whole b while doing the transposition?

It seems like this is a fairly ordinary operation so I'm confused that it does not work.

Solution

Your interpretation is not far off the mark: numba has four candidates to multiply C and F layout arrays, and gives you the details, why each one is not chosen in the end. The latter two are dismissed because an argument is missing, so they are clearly intended for another call signature. The former two are dismissed because something does not work internally:

LoweringError: Failed in nopython mode pipeline (step: native lowering)
          scipy 0.16+ is required for linear algebra

While the first line is pretty cryptic, the second is still part of the error message and gives a pretty good hint. Install scipy manually and it should work.

On a side note: that's basically a one-liner of numpy functions, which should perform pretty well on a single CPU core as there's not much Python overhead to eliminate for numba. It of course depends on what else you want to do, but don't expect this single piece to get a significant boost.

Answered By - Matthias Huschle