Pallas autotuning

[lengstrom]

I could not fine an interface for autotuning (as in Triton: https://triton-lang.org/main/python-api/generated/triton.autotune.html) in Pallas. Is there currently a way of doing this?

If not, is there a way to access triton autotuning from Pallas?

[justinjfu]

Hi,

We don't have a system for autotuning yet (there's an old prototype here: jax-ml/jax-triton#108). I can't promise a release date in the near future but it's on our list of features to implement. In the meantime, we normally write a separate script to benchmark and perform a grid search over parameters and save the best ones manually.