xla

e1e2357b - Rename prelu to _prelu_kernel and update diag_embed op to functioanlization pass instead of manual lowering