From d2337ee7ef5c20c91b44c49f9637eebd961b8468 Mon Sep 17 00:00:00 2001
From: Sharad Vikram <sharadmv@google.com>
Date: Fri, 3 Jan 2025 15:29:32 -0800
Subject: [PATCH] [Mosaic TPU] Add some elementwise canonicalizations

PiperOrigin-RevId: 711869914
---
 jaxlib/mosaic/dialect/tpu/transforms/canonicalize_mosaic.cc | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/jaxlib/mosaic/dialect/tpu/transforms/canonicalize_mosaic.cc b/jaxlib/mosaic/dialect/tpu/transforms/canonicalize_mosaic.cc
index 567838a8cab1..44ecb600bbba 100644
--- a/jaxlib/mosaic/dialect/tpu/transforms/canonicalize_mosaic.cc
+++ b/jaxlib/mosaic/dialect/tpu/transforms/canonicalize_mosaic.cc
@@ -552,7 +552,10 @@ const llvm::StringSet<> &elementwise_convertible_ops() {
                                           arith::SubFOp::getOperationName(),
                                           arith::MaximumFOp::getOperationName(),
                                           arith::MinimumFOp::getOperationName(),
-                                          math::PowFOp::getOperationName()};
+                                          math::PowFOp::getOperationName(),
+                                          math::TanhOp::getOperationName(),
+                                          math::ExpOp::getOperationName(),
+                                          math::LogOp::getOperationName()};
   return *ops;
 }