ASEM000 / kernex

Stencil computations in JAX
MIT License
66 stars 3 forks source link

add argnum/argname #6

Open ASEM000 opened 1 year ago

ASEM000 commented 1 year ago

kernex operate on views of the first function array argument, it is useful to add argnum/argname to select which argument to operate on in a similar fashion in jax transformation