Numba JIT TypingErrors Thrown on Multipack Functions

[fabianlim]

Multipack requires runing through the dataset and computing some statistics at the beginning at every epoch. It was reported by @anhuong that when setting the loglevel in DEBUG that we observe some TypingErrors during the JIT compilation

It is unclear if there is a severe effect, since the function still executes correctly. It could be that it internally fallsback to the non-compiled version. But regardless this is a numba JIT error

from fms_acceleration_aadp.multipack_sampler import allocate
import numpy as np
import logging

logging.basicConfig(level=logging.DEBUG)
logger = logging.getLogger()

if __name__ == '__main__':

    logger.debug("this is a test")
    lengths = np.random.randint(500, 1000, 1000)

    batches, total_used, total_slots = allocate(
        lengths, np.cumsum(lengths),
        rank=0, c=5000, n=8, padding=False
    )
    print ("batches", batches)

The TypingError happens in a few places

Traceback (most recent call last):
  File "/workspace/fms/run-benches/lib/python3.10/site-packages/numba/core/typeinfer.py", line 160, in propagate
    constraint(typeinfer)
  File "/workspace/fms/run-benches/lib/python3.10/site-packages/numba/core/typeinfer.py", line 581, in __call__
    fnty = typevars[self.func].getone()
  File "/workspace/fms/run-benches/lib/python3.10/site-packages/numba/core/typeinfer.py", line 128, in getone
    raise TypingError("Undecided type {}".format(self))
numba.core.errors.TypingError: Undecided type $152load_method.20 := <undecided>
During: resolving caller type: $152load_method.20
During: typing of call at /workspace/fms-acceleration/plugins/attention-and-distributed-packing/src/fms_acceleration_aadp/multipack_sampler.py (331)

Traceback (most recent call last):
  File "/workspace/fms/run-benches/lib/python3.10/site-packages/numba/core/typeinfer.py", line 160, in propagate
    constraint(typeinfer)
  File "/workspace/fms/run-benches/lib/python3.10/site-packages/numba/core/typeinfer.py", line 581, in __call__
    fnty = typevars[self.func].getone()
  File "/workspace/fms/run-benches/lib/python3.10/site-packages/numba/core/typeinfer.py", line 128, in getone
    raise TypingError("Undecided type {}".format(self))
numba.core.errors.TypingError: Undecided type $108load_method.12 := <undecided>
During: resolving caller type: $108load_method.12
During: typing of call at /workspace/fms-acceleration/plugins/attention-and-distributed-packing/src/fms_acceleration_aadp/multipack_sampler.py (299)

Traceback (most recent call last):
  File "/workspace/fms/run-benches/lib/python3.10/site-packages/numba/core/typeinfer.py", line 160, in propagate
    constraint(typeinfer)
  File "/workspace/fms/run-benches/lib/python3.10/site-packages/numba/core/typeinfer.py", line 581, in __call__
    fnty = typevars[self.func].getone()
  File "/workspace/fms/run-benches/lib/python3.10/site-packages/numba/core/typeinfer.py", line 128, in getone
    raise TypingError("Undecided type {}".format(self))
numba.core.errors.TypingError: Undecided type $108load_method.12 := <undecided>
During: resolving caller type: $108load_method.12

log.txt