FEAT - Use @njit(cache=True)

[PascalCarrivain]

Use caching option from Numba

I explore a little bit the skglm source code and I realized you are using Numba decorator @nijt.
I was wondering if it makes senses to switch to @nijt(cache=True).
Indeed, according to Numba documentation caching compiled functions reduces the future compilation time.

[mathurinm]

Thanks for the pointer @PascalCarrivain, I was not aware of this feature.

It seems to help a lot (not for the first compilation, but for the subsequent calls), on a very CPU bound problem. From the first run to the second, I change only the value of cache in the following snippet)

In [1]: %run numba_cache.py
0.5166642830008641
0.10939704200063716
0.11151579199940898
0.11363932000131172
0.10976769300032174

In [2]:                                                                         
Do you really want to exit ([y]/n)? 
(base) ➜  scripts git:(main) ✗ ipython
Python 3.10.12 | packaged by conda-forge | (main, Jun 23 2023, 22:40:32) [GCC 12.3.0]
Type 'copyright', 'credits' or 'license' for more information
IPython 8.6.0 -- An enhanced Interactive Python. Type '?' for help.

In [1]: %run numba_cache.py
0.529917010000645
0.004792376999830594
0.004180643998552114
0.004381413000373868
0.004056714000398642

import numpy as np
import time
from numba import njit

a = np.arange(10_000)

for i in range(5):
    def my_sum(a):
        acc = 0
        for val in a:
            acc += val
        return acc

    t0 = time.perf_counter()
    njit(my_sum, cache=True)(a)
    t1 = time.perf_counter()
    print(t1 - t0)

Can you try to test the impact of using cache=True in our codebase on a real life skglm problem, ie fitting an estimator on a simple problem ?

[PascalCarrivain]

@mathurinm Yes, I will do it late this year or early next year.

[mathurinm]

@PascalCarrivain do you know if this can help the first compilation too ?

[PascalCarrivain]

I do not see a huge difference for the first compilation (at least on my projects).

[PascalCarrivain]

Do you have a specific example you would like to benchmark with and without cache=True?

[mathurinm]

Something like fitting a lasso on a dataset of reasonable size should be relevant Il mar 22 apr 2025, 21:54 Pascal Carrivain ***@***.***> ha scritto:

…

Do you have a specific example you would like to benchmark with and without cache=True? — Reply to this email directly, view it on GitHub <#215 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/ACETTQUR6WKB625BR6YM75L222M6DAVCNFSM6AAAAAB3UR2X5SVHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMZDQMRSGMZTKMZQGI> . You are receiving this because you were mentioned.Message ID: ***@***.***> *PascalCarrivain* left a comment (scikit-learn-contrib/skglm#215) <#215 (comment)> Do you have a specific example you would like to benchmark with and without cache=True? — Reply to this email directly, view it on GitHub <#215 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/ACETTQUR6WKB625BR6YM75L222M6DAVCNFSM6AAAAAB3UR2X5SVHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMZDQMRSGMZTKMZQGI> . You are receiving this because you were mentioned.Message ID: ***@***.***>