Avoid specialization explosion for heterogeneous tuples as arguments

Numba Support: How do I do ...?
1

Hello everyone,

How can we avoid the explosion of multiple specializations when compiling functions that take a heterogeneous tuple of arrays as an argument (perhaps using intrinsics)?

Here’s an example where the function receives a tuple of arrays.

  • The allowed array types are int64 and float64.
  • The number of arrays ranges from 1 to 5.
  • The order of array types within the tuple is random.

This setup results in 62 specializations.
Can we somehow achieve less specializations without restricting the input arguments?

import numpy as np
from numba import jit, literal_unroll
from itertools import product

@jit
def foo(tup):
    size = 0
    for a in literal_unroll(tup):
        size += a.size
    return size

# arrays can be flotat64 or int64
arr_i64, arr_f64 = np.arange(3), np.arange(3.)
arrays = [arr_f64, arr_i64]

# Generate all possible combinations of 1 to 5 arrays
combinations = []
for r in range(1, 6):
    combinations.extend(product(arrays, repeat=r))

for tup in combinations:
    foo(tup)

print("Number of signatures for heterogeneous tuple of arrays")
print("with 1-2 dtypes and 1-5 arrays:")
print(f"count: {len(foo.signatures)}")
# count: 62