xla

f0881b5a - Use f32 scratch for output so we only need to transfer output with desired dtype back to HBM. (#8924)