Use of `Enum` and `numba`

[felixhekhorn]

@kamillaurent reports:

2026-05-28 03:20:21,201 eko.evolution_operator/INFO: Evolution: computing operators - 193/196 took: 0.909104 s
2026-05-28 03:20:22,105 eko.evolution_operator/INFO: Evolution: computing operators - 194/196 took: 0.903717 s
2026-05-28 03:20:23,390 eko.evolution_operator/INFO: Evolution: computing operators - 195/196 took: 0.897556 s
2026-05-28 03:20:24,109 eko.evolution_operator/INFO: Evolution: computing operators - 196/196 took: 0.717515 s
2026-05-28 03:20:25,765 eko.evolution_operator/INFO: Evolution: Total time 4963.983533 s
[... after some time ...]
Traceback (most recent call last):
  File "/data/theorie/klaurent/miniconda3/envs/nnpdf-dev/bin/evolven3fit", line 7, in <module>
    sys.exit(main())
             ^^^^^^
  File "/data/theorie/klaurent/nnpdf/n3fit/src/n3fit/scripts/evolven3fit.py", line 187, in main
    solve(tcard, opcard, args.dump)
  File "/data/theorie/klaurent/miniconda3/envs/nnpdf-dev/lib/python3.12/site-packages/eko/runner/managed.py", line 29, in solve
    eko.parts[recipe] = parts.evolve(eko, recipe)
                        ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/data/theorie/klaurent/miniconda3/envs/nnpdf-dev/lib/python3.12/site-packages/eko/runner/parts.py", line 76, in evolve
    op.compute()
  File "/data/theorie/klaurent/miniconda3/envs/nnpdf-dev/lib/python3.12/site-packages/eko/evolution_operator/__init__.py", line 981, in compute
    self.integrate()
  File "/data/theorie/klaurent/miniconda3/envs/nnpdf-dev/lib/python3.12/site-packages/eko/evolution_operator/__init__.py", line 1001, in integrate
    for j, row in enumerate(res):
  File "/data/theorie/klaurent/miniconda3/envs/nnpdf-dev/lib/python3.12/site-packages/eko/evolution_operator/__init__.py", line 902, in run_op_integration
    res = integrate.quad(
          ^^^^^^^^^^^^^^^
  File "/data/theorie/klaurent/miniconda3/envs/nnpdf-dev/lib/python3.12/site-packages/scipy/integrate/_quadpack_py.py", line 479, in quad
    retval = _quad(func, a, b, args, full_output, epsabs, epsrel, limit,
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/data/theorie/klaurent/miniconda3/envs/nnpdf-dev/lib/python3.12/site-packages/scipy/integrate/_quadpack_py.py", line 626, in _quad
    return _quadpack._qagse(func,a,b,args,full_output,epsabs,epsrel,limit)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/data/theorie/klaurent/miniconda3/envs/nnpdf-dev/lib/python3.12/site-packages/numba/core/dispatcher.py", line 687, in typeof_pyval
    tp = typeof(val, Purpose.argument)
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/data/theorie/klaurent/miniconda3/envs/nnpdf-dev/lib/python3.12/site-packages/numba/core/typing/typeof.py", line 33, in typeof
    ty = typeof_impl(val, c)
         ^^^^^^^^^^^^^^^^^^^
  File "/data/theorie/klaurent/miniconda3/envs/nnpdf-dev/lib/python3.12/functools.py", line 909, in wrapper
    return dispatch(args[0].__class__)(*args, **kw)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/data/theorie/klaurent/miniconda3/envs/nnpdf-dev/lib/python3.12/site-packages/numba/core/typing/typeof.py", line 230, in _typeof_enum
    clsty = typeof_impl(type(val), c)
            ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/data/theorie/klaurent/miniconda3/envs/nnpdf-dev/lib/python3.12/functools.py", line 909, in wrapper
    return dispatch(args[0].__class__)(*args, **kw)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/data/theorie/klaurent/miniconda3/envs/nnpdf-dev/lib/python3.12/site-packages/numba/core/typing/typeof.py", line 249, in _typeof_enum_class
    return typecls(cls, dtypes.pop())
           ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/data/theorie/klaurent/miniconda3/envs/nnpdf-dev/lib/python3.12/site-packages/numba/core/types/abstract.py", line 68, in __call__
    return cls._intern(inst)
           ^^^^^^^^^^^^^^^^^
  File "/data/theorie/klaurent/miniconda3/envs/nnpdf-dev/lib/python3.12/site-packages/numba/core/types/abstract.py", line 57, in _intern
    inst._code = _autoincr()
                 ^^^^^^^^^^^
  File "/data/theorie/klaurent/miniconda3/envs/nnpdf-dev/lib/python3.12/site-packages/numba/core/types/abstract.py", line 23, in _autoincr
    assert n < 2 ** 32, "Limited to 4 billion types"
AssertionError: Limited to 4 billion types

The problem seems to be (according to Claude run by @juanrojochacon )

The trigger is `_typeof_enum` → `_typeof_enum_class` in the chain: every time EKO calls a Numba-jitted integrand through `scipy.integrate.quad`, Numba inspects the arguments and, for an `Enum` argument, **creates a brand-new `EnumClass` type and interns it**. Numba does not cache enum types by identity properly here, so each call leaks a fresh type into the global registry. EKO's `run_op_integration` runs `quad` an enormous number of times (per grid point, per operator element, per Q² step), so the counter climbs until it overflows 2³².

and a supposed solutions (suggested by the same pair of "people")

The cleanest fix is to stop passing the enum into the jitted function and pass its `.value` (a plain int) instead, converting back inside or branching outside the hot loop

[felixhekhorn]

If that would be true (which I don't know yet) the effected variables would be for Operator

eko/src/eko/evolution_operator/__init__.py

Line 833 in 9dc7fa5

method=self.ev_method,

eko/src/eko/evolution_operator/__init__.py

Line 846 in 9dc7fa5

sv_mode=self.sv_mode,

and for OME

eko/src/eko/evolution_operator/operator_matrix_element.py

Line 331 in 217c326

sv_mode=self.sv_mode,

eko/src/eko/evolution_operator/operator_matrix_element.py

Line 333 in 217c326

backward_method=self.backward_method,

I think

Casting twice, 1) to int before quad_ker* + 2) back to Enum inside quad_ker, might be the short-term way out

[kamillaurent]

Since when I install nnpdf in a clean conda environment eko is a dependency, I am now applying the suggested changes into the the dependency in my conda env.

I am trying to re-compute the PDF EKO for theory 41000100 and see if this work. If this change (passing integers values into numba-jitted quad kernels instead of Enum) solves the issue, I can open a PR for eko with the small fix.

Do you think this is a good plan or there is a more robust way of proceeding @felixhekhorn?

[felixhekhorn]

First, we need to understand if Claude is right - and that I can not confirm so far ...

My test script:

# cat numba_enum.py
import numba as nb
from scipy.integrate import quad
from enum import IntEnum, auto

class Test(IntEnum):
    A = auto()
    B = auto()

@nb.njit
def test(x,t):
    if t == Test.A:
        return 0.
    return 1.
    
print(quad(test,0.,1.,args=(Test.A)))
print(quad(test,0.,1.,args=(Test.B)))
print(quad(test,0.,1.,args=(Test.A)))

Running $ NUMBA_DEBUG=1 python numba_enum.py > numba_enum.log writes a whole mess of all sorts of complicated stuff out, which I don't understand ... However, if I search inside that mess for "enum" I get the following stuff:

[... 315 lines ...]
---------------------------------IR DUMP: test----------------------------------
label 0:
    x = arg(0, name=x)                       ['x']
    t = arg(1, name=t)                       ['t']
    $4load_global.1 = global(Test: <enum 'Test'>) ['$4load_global.1']
    $6load_attr.2 = getattr(value=$4load_global.1, attr=A) ['$4load_global.1', '$6load_attr.2']
    $8compare_op.3 = t == $6load_attr.2      ['$6load_attr.2', '$8compare_op.3', 't']
    $bool10 = global(bool: <class 'bool'>)   ['$bool10']
    $10pred = call $bool10($8compare_op.3, func=$bool10, args=(Var($8compare_op.3, numba_enum.py:11),), kws=(), vararg=None, varkwarg=None, target=None) ['$10pred', '$8compare_op.3', '$bool10']
    branch $10pred, 12, 16                   ['$10pred']
label 12:
    $const12.0.1 = const(float, 0.0)         ['$const12.0.1']
    $14return_value.1 = cast(value=$const12.0.1) ['$14return_value.1', '$const12.0.1']
    return $14return_value.1                 ['$14return_value.1']
label 16:
    $const16.0.2 = const(float, 1.0)         ['$const16.0.2']
    $18return_value.1 = cast(value=$const16.0.2) ['$18return_value.1', '$const16.0.2']
    return $18return_value.1                 ['$18return_value.1']

-------------------------------SSA IR DUMP: test--------------------------------
label 0:
    x = arg(0, name=x)                       ['x']
    t = arg(1, name=t)                       ['t']
    $4load_global.1 = global(Test: <enum 'Test'>) ['$4load_global.1']
    $6load_attr.2 = getattr(value=$4load_global.1, attr=A) ['$4load_global.1', '$6load_attr.2']
    $8compare_op.3 = t == $6load_attr.2      ['$6load_attr.2', '$8compare_op.3', 't']
    $bool10 = global(bool: <class 'bool'>)   ['$bool10']
    $10pred = call $bool10($8compare_op.3, func=$bool10, args=(Var($8compare_op.3, numba_enum.py:11),), kws=(), vararg=None, varkwarg=None, target=None) ['$10pred', '$8compare_op.3', '$bool10']
    branch $10pred, 12, 16                   ['$10pred']
label 12:
    $const12.0.1 = const(float, 0.0)         ['$const12.0.1']
    $14return_value.1 = cast(value=$const12.0.1) ['$14return_value.1', '$const12.0.1']
    return $14return_value.1                 ['$14return_value.1']
label 16:
    $const16.0.2 = const(float, 1.0)         ['$const16.0.2']
    $18return_value.1 = cast(value=$const16.0.2) ['$18return_value.1', '$const16.0.2']
    return $18return_value.1                 ['$18return_value.1']

-----------------------------------propagate------------------------------------
---- type variables ----
[$10pred := bool,
 $14return_value.1 := float64,
 $18return_value.1 := float64,
 $4load_global.1 := IntEnum class<int64>(Test),
 $6load_attr.2 := IntEnum<int64>(Test),
 $8compare_op.3 := bool,
 $bool10 := Function(<class 'bool'>),
 $const12.0.1 := float64,
 $const16.0.2 := float64,
 arg.t := IntEnum<int64>(Test),
 arg.x := float64,
 t := IntEnum<int64>(Test),
 x := float64]
-----------------------------------propagate------------------------------------
---- type variables ----
[$10pred := bool,
 $14return_value.1 := float64,
 $18return_value.1 := float64,
 $4load_global.1 := IntEnum class<int64>(Test),
 $6load_attr.2 := IntEnum<int64>(Test),
 $8compare_op.3 := bool,
 $bool10 := Function(<class 'bool'>),
 $const12.0.1 := float64,
 $const16.0.2 := float64,
 arg.t := IntEnum<int64>(Test),
 arg.x := float64,
 t := IntEnum<int64>(Test),
 x := float64]
---------------------------------Variable types---------------------------------
{'$10pred': bool,
 '$14return_value.1': float64,
 '$18return_value.1': float64,
 '$4load_global.1': IntEnum class<int64>(Test),
 '$6load_attr.2': IntEnum<int64>(Test),
 '$8compare_op.3': bool,
 '$bool10': Function(<class 'bool'>),
 '$const12.0.1': float64,
 '$const16.0.2': float64,
 'arg.t': IntEnum<int64>(Test),
 'arg.x': float64,
 't': IntEnum<int64>(Test),
 'x': float64}
----------------------------------Return type-----------------------------------
float64
[... 606 more lines ...]

if we then guess a bit in the wild and think that $4load_global.1 is a global registry, we see this does not get registered twice. The pasted snippet contains all appearances of "enum". A simpler version of the snippet without scipy in the middle, yields basically the same stuff. This is by no means any proof or disproof.

Anyways, I'd like to be sure that this is the problem before we do anything about it. And in the numba repository there is no such problem reported (yet). If the conjecture is true, that would be a numba bug and while we wait for them for a fix, we can then consider apply a fix locally.

[lol782]

hey @felixhekhorn could you please elaborate on the Enum/numba issue? Happy to work on it once I understand the specific problem.

[felixhekhorn]

I'm not sure myself what the problem is or even if there is one (or if it was just a local problem), so this is not an easy issue I would say