xla

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