Simplifying lowering and abstract evaluation rules

flatironinstitute / jax-finufft

JAX bindings to the Flatiron Institute Non-uniform Fast Fourier Transform (FINUFFT) library

Apache License 2.0

77 stars 2 forks source link

Simplifying lowering and abstract evaluation rules #89

Closed dfm closed 4 months ago

dfm commented 4 months ago

With relatively recent versions of JAX we can remove some of the boilerplate in the MLIR lowering rules.

dfm commented 4 months ago

@lgarrison — Good point! It looks like the oldest version of JAX that this supports is 0.4.20 from Nov 2023. This would make us somewhat bleeding edge, but I'm not too concerned about that. What do you think?

lgarrison commented 4 months ago

I think that sounds fine, we already have a release that works with older JAX on PyPI. LGTM!