[math] Update CustomOpByNumba to support JAX version >= 0.4.24 (#669)
#1171
| Job | Run time |
|---|---|
| 1m 30s | |
| 56s | |
| 1m 1s | |
| 2m 14s | |
| 51s | |
| 55s | |
| 2m 33s | |
| 1m 19s | |
| 57s | |
| 12m 16s |