Issue with numba: "cannot unify array" unless I convert scalars to array

I have found a solution, but it's far from elegant, and I am hoping there is a better one.

To recap, I needed a function which:

• works with numba
• works with both scalars and arrays
• returns scalar (not a one-sized array) when the inputs are scalars, and arrays when the inputs are arrays

I have tried the following, and found option 2 to be the fastest.

1. my_fun_optimised_1: a function which, without numba, determines whether the inputs are scalar or not, and then calls, accordingly, a sub-function for the scalar case and one for the arrays. Both sub-functions run with numba, but take forever. I guess this is because numba must be re-initialised at each call of the main function.
2. my_fun_optimised_2: similar to the above, except the scalar and array functions, both running with numba, are main functions and not subfunctions. Much much faster.
3. my_fun_non_opt_no_numba : a function which runs without numba.

The results are:

+-------------------------+----------------------------+-----------------------------+
|        Function         | Array: time vs the fastest | Scalar: time vs the fastest |
+-------------------------+----------------------------+-----------------------------+
| optimised numba 1       |                     54,403 |                      42,961 |
| optimised numba 2       |                          1 |                           1 |
| non-optimised, no numba |                      3.409 |                     4.53892 |
+-------------------------+----------------------------+-----------------------------+

What this means is that, on my PC, the non-optimised, no numba code takes 4.5 times longer than "optimsied numba 2" to run on scalars and 3.4 times longer for arrays. The "optimised numba 1" is not optimsied at all and takes an insane amount of time.

I hope all of this can be of use to other people.

PS I am very well familiar with pitfalls of premature optimisation. I am only doing this because I have a specific case where 60% of the time is spent doing a similar (but not identical) calculation to the one shown here.

The code to time the functions is:

import numpy as np
import numba
import timeit
import pandas as pd

def my_fun_optimised_1(a,b,c):

    @numba.jit
    def my_fun_vectorised(a,b,c):
        return np.where(a == 0, np.nan,
                        np.where(b ==0 , 0 , b*a**3 + a*b**3 + a*b*c**3)
                        )
    
    @numba.jit
    def my_fun_scalar(a,b,c):
        if a ==0:
            return np.nan
        elif b == 0:
            return np.nan
        else:
            return b*a**3 + a*b**3 + a*b*c**3
        
    if np.isscalar(a) and np.isscalar(b) and np.isscalar(c):
        return my_fun_scalar(a,b,c)
    else:        
        return my_fun_vectorised(a,b,c)
    

def my_fun_optimised_2(a,b,c):
    if np.isscalar(a) and np.isscalar(b) and np.isscalar(c):
        return fun_2_scalar(a,b,c)
    else:        
        return fun_2_vectorised(a,b,c)
    
@numba.jit
def fun_2_scalar(a,b,c):
    if a ==0:
        return np.nan
    elif b == 0:
        return np.nan
    else:
        return b*a**3 + a*b**3 + a*b*c**3

@numba.jit
def fun_2_vectorised(a,b,c):
    return np.where(a == 0, np.nan,
                    np.where(b ==0 , 0 , b*a**3 + a*b**3 + a*b*c**3)
                    )
    
    
    
def my_fun_non_opt_no_numba(a,b,c):
    # multipl by 1 converts from array to scalar
    return  1 * np.where(a == 0, np.nan,
                    np.where(b ==0 , 0 , b*a**3 + a*b**3 + a*b*c**3) 
                    )

# I couldn't get this to work with Numba
#@numba.jit
def my_fun_non_opt_numba(a,b,c):
    a = np.array([a])
    b = np.array([b])
    c = np.array([c])
    out = np.where(a == 0, np.nan,
                    np.where(b ==0 , 0 , b*a**3 + a*b**3 + a*b*c**3) 
                    )
    return out





r = 4
n = int(100)

a = 3
b = 4
c = 5

x = my_fun_optimised_2(a,b,c)

t_scalar_opt_numba_1 = timeit.Timer("my_fun_optimised_1(a,b,c) " ,  globals = globals() ).repeat(repeat = r, number = n)
t_scalar_opt_numba_2 = timeit.Timer("my_fun_optimised_2(a,b,c) " ,  globals = globals() ).repeat(repeat = r, number = n)
t_scalar_non_opt_no_numba = timeit.Timer("my_fun_non_opt_no_numba(a,b,c) " ,  globals = globals() ).repeat(repeat = r, number = n)

resdf_scalar = pd.DataFrame(index = ['min time'])
resdf_scalar['optimised numba 1'] = [min(t_scalar_opt_numba_1)]
resdf_scalar['optimised numba 2'] = [min(t_scalar_opt_numba_2)]
resdf_scalar['non-optimised, no numba'] = [min(t_scalar_non_opt_no_numba)]
# the docs explain why we should take the min and not the avg
resdf_scalar = resdf_scalar.transpose()
resdf_scalar['diff vs fastest'] = (resdf_scalar / resdf_scalar.min() )



a = np.arange(3,13)
b = np.arange(0,10)
c = np.arange(20,30)

t_array_opt_numba_1 = timeit.Timer("my_fun_optimised_1(a,b,c) " ,  globals = globals() ).repeat(repeat = r, number = n)
t_array_opt_numba_2 = timeit.Timer("my_fun_optimised_2(a,b,c) " ,  globals = globals() ).repeat(repeat = r, number = n)
t_array_non_opt_no_numba = timeit.Timer("my_fun_non_opt_no_numba(a,b,c) " ,  globals = globals() ).repeat(repeat = r, number = n)

resdf_array = pd.DataFrame(index = ['min time'])
resdf_array['optimised numba 1'] = [min(t_array_opt_numba_1)]
resdf_array['optimised numba 2'] = [min(t_array_opt_numba_2)]
resdf_array['non-optimised, no numba'] = [min(t_array_non_opt_no_numba)]
# the docs explain why we should take the min and not the avg
resdf_array = resdf_array.transpose()
resdf_array['diff vs fastest'] = (resdf_array / resdf_array.min() )