kernex (pypi): Stencil computations in JAX.
Keywords: python|machine-learning|pytorch|jax|kernel|stencil