How to pass a part of args within closure and others when calling inner function?

[Roy-Kid]

Hi, I have a question that I can not fix for a long time. The hierarchy of the project is:

jaxtest
    src
        * main.py
    * test.py

in the main.py, I write a closure:

def outfunc(criterion):
    criterion = criterion + 1  # do something
    def newfunc(a, b):
        if criterion % 2 == 0:  # condition
            return a
        else:
            return b
    return jit(newfunc)

and in the test.py

func = outfunc(np.array(2))
print(func(0, 1))

What I want to is when calling outfunc, I need to pass some args, and those args determine the behavior of newfunc. If there is an if condition statement in the jitted function, we must indicate args as static_argnums or names. But if I write like this, return jit(newfunc, static_argnames='criterion') still, ConcretizationTypeError raises.

Thanks in advance if someone can help me with this.