Using closure to reuse code

Hi, I’m trying to compose njit-ed functions that use closure values.
Is there any way to make this pattern work, while keeping h defined outside f?

def h(g):
    return g(1)
def f(a,b):
    def g(i):
        return b[i]
    return h(g)
TypingError: Failed in nopython mode pipeline (step: convert make_function into JIT functions)
Cannot capture the non-constant value associated with variable 'b' in a function that will escape.

File "<ipython-input-179-3da212f9c5bf", line 6:
def f(a,b):
    def g(i):

‘b’ here is not mutated thorough the algorithm, and ‘g’ could be inlined.

I arrived at an initial solution using @jitclass to emulate functools.partial.
I tried to improve the interface with dynamic class generation and the use of @generated_jit.
Example below:

import pyrsistent
from numba.experimental import jitclass

from REDACTED.util.indented_fstring import deindent
from weakref import WeakValueDictionary
def numba_closure_func(name,func,const={},variable=()):
    global _numba_closure_cache
    key = pyrsistent.freeze((func,const,variable))
    if key in _numba_closure_func_cache:
        return _numba_closure_func_cache[key]
    spec = [(k,v) for k,v in const.items()]
    call_variable_args = ','+','.join(variable) if variable else ''
    func_variable_args = ','+','.join(variable) if variable else ''
    func_args = ','.join(f"self.{k}" for k in const.keys())+func_variable_args
    cls_def =deindent(f"""\
        class {name}:
            def __init__(self{','+','.join(const.keys()) if const else ''}):
                {';'.join(f'self.{k}={k}' for k in const.keys()) if const else 'pass'}
            def call(self{call_variable_args}):
                return func({func_args})
    lg = dict(func=func)
    ld = dict()
    kls = ld[name]
    klsj = jitclass(spec)(kls)
    return klsj

import numpy as np

def G_func(b,i):
    return b[i]
def h(g):

def g_func(b,i):
    return b[i]

def f(a,b):
    G = numba_closure_func('G',g_func,const={'b':b},variable=('i',))
    def _f(a,b):
        g = G(b)
        return h(g)
    return _f
#print of spec and class def
[('b', array(float64, 1d, C))]
class G:
    def __init__(self,b):
    def call(self,i):
        return func(self.b,i)
# correct result

hi there! was this not an option for you

    def h(g, b):
        return g(1, b)
    def f(a,b):
        def g(i, b):
            return b[i]
        return h(g, b)

That would carry g’s signature into h’s, making h less general; I would have to have a h written for each possible signature (signature as in number of arguments) of g.