[math] Update `CustomOpByNumba` to support JAX version >= 0.4.24

brainpy / BrainPy

Brain Dynamics Programming in Python

https://brainpy.readthedocs.io/

GNU General Public License v3.0

514 stars 93 forks source link

Closed Routhleck closed 4 months ago

Routhleck commented 4 months ago

Update CustomOpByNumba to support JAX version >= 0.4.24 by implementing mlir register rules