[Mosaic TPU] Only canonicalize the dtype of transposes if they use the XLU

We should never change the dtypes of transposes that only deal with untiled dimensions,
for example.

PiperOrigin-RevId: 815013892

Author: apaszke

jaxlib/mosaic/dialect/tpu/transforms/canonicalize_mosaic.cc

@@ -1799,9 +1799,11 @@ FailureOr<Value> canonicalize_transpose(const CanonicalizeContext &ctx,
     return res;
   };
 
+  bool uses_xlu = !op.getPermutation().empty() &&
+                  op.getPermutation().back() != op.getPermutation().size() - 1;
   // TODO(b/448848595): Enable 8-bit transposes on generation 7.
   if (element_type.getIntOrFloatBitWidth() == 8 && ctx.compatibility_mode &&
-      ctx.hardware_generation > 3) {
+      ctx.hardware_generation > 3 && uses_xlu) {
     VectorType input_vty_int = VectorType::get(
         input_vty.getShape(),
         builder.getIntegerType(input_vty.getElementTypeBitWidth()));