Отправка numba по типу

Вам не следует использовать isinstance в такой функции JIT, а вместо этого используйте @overload ( @generated_jit был старым устаревшим способом сделать это), который специально создан для этой цели. Это позволяет Numba генерировать код быстрее, поскольку для каждого случая компилируется только часть функции, а не весь случай для каждой специализации. Более того, isinstance является экспериментальным, как указано Numba в предупреждении при выполнении вашего первого кода (предупреждения выдаются, чтобы пользователи могли их прочитать ;)).

Используя новый метод @overload

Начиная с Numba 0.59, вместо него необходимо использовать перегрузку:

import numba as nb
import numpy as np

def test_dispatch_scalar(X, indices):
    ref_pos = np.empty(3, np.float64)
    ref_pos[:] = X[:, indices]
    return ref_pos

def test_dispatch_vector(X, indices):
    ref_pos = np.empty((3, len(indices)), np.float64)
    ref_pos[:, :] = X[:, indices]
    return ref_pos

# Pure-python fallback implementation
def test_dispatch_impl(X, indices):
    if isinstance(indices, (int, np.integer)):
        return test_dispatch_scalar(X, indices)
    elif isinstance(indices, np.ndarray) and indices.ndim == 1 and np.issubdtype(indices.dtype, np.integer):
        return test_dispatch_vector(X, indices)
    else:
        assert False # Unsupported

# Numba-specific overload
@nb.extending.overload(test_dispatch_impl)
def test_dispatch_impl_overload(X, indices):
    if isinstance(indices, nb.types.Integer):
        return test_dispatch_scalar
    elif isinstance(indices, nb.types.Array) and indices.ndim == 1 and isinstance(indices.dtype, nb.types.Integer):
        return test_dispatch_vector
    else:
        assert False # Unsupported

@nb.njit
def test_dispatch(X, indices):
    return test_dispatch_impl(X, indices)

Старое устаревшее решение

Вот пример рассуждений об универсальных типах:

import numba as nb
import numpy as np

@nb.generated_jit(nopython=True)
def test_dispatch(X, indices):
    if isinstance(indices, nb.types.Integer):
        def test_dispatch_scalar(X, indices):
            ref_pos = np.empty(3, np.float64)
            ref_pos[:] = X[:, indices]
            return ref_pos
        return test_dispatch_scalar
    elif isinstance(indices, nb.types.Array) and indices.ndim == 1 and isinstance(indices.dtype, nb.types.Integer):
        def test_dispatch_vector(X, indices):
            ref_pos = np.empty((3, len(indices)), np.float64)
            ref_pos[:, :] = X[:, indices]
            return ref_pos
        return test_dispatch_vector
    else:
        assert False # Unsupported

Вот пример рассуждений о конкретных типах:

import numba as nb
import numpy as np

@nb.generated_jit(nopython=True)
def test_dispatch(X, indices):
    if indices == nb.types.int64:
        def test_dispatch_scalar(X, indices):
            ref_pos = np.empty(3, np.float64)
            ref_pos[:] = X[:, indices]
            return ref_pos
        return test_dispatch_scalar
    elif isinstance(indices, nb.types.Array) and indices.ndim == 1 and indices.dtype == nb.types.int64:
        def test_dispatch_vector(X, indices):
            ref_pos = np.empty((3, len(indices)), np.float64)
            ref_pos[:, :] = X[:, indices]
            return ref_pos
        return test_dispatch_vector
    else:
        assert False # Unsupported

Запрос конкретных 64-битных целых чисел может быть слишком ограничительным, поэтому я советую вам смешивать тесты общего типа и тесты конкретных типов. По той же причине вам следует избегать прямого тестирования массивов определенного типа просто потому, что они часто могут быть смежными или нет или могут содержать типы элементов, совместимые с вашей функцией.

Обратите внимание, что общие функции JIT предназначены для создания функций, которые компилируются отдельно относительно целевого типа входных данных (а не значений).