Does JAX support post training quantization/downcasting?

[alexbeloi]

ATT

Ideally would like to be able to try fp16 and bf16.

I'm aware of some XLA flags like XLA_USE_BF16 and XLA_DOWNCAST_BF16 but they do nothing in my case, I believe they're either TorchXLA and/or TPU specific.

If not, would it be possible to do this manually by say AOT lowering and modifying the HLO before compiling it?