Fix tpu (#1054)

* Fix tpu

* remove llvm operator on tpu

* fix

* fix

Author: wsmoses

src/Compiler.jl

@@ -1189,13 +1189,22 @@ function compile_mlir!(
     # Add a `donated` attr to the function arguments. This doesn't affect XLA, but lets us
     # check which arguments were donated.
     preserved_args_idx = last.(preserved_args)
+    if backend != "tpu"
     for (i, arg) in enumerate(linear_args)
         if i ∉ preserved_args_idx
             MLIR.API.mlirFuncSetArgAttr(
                 func3, i - 1, "reactant.donated", MLIR.IR.UnitAttribute()
             )
         end
     end
+    else
+        for op in collect(MLIR.IR.OperationIterator(MLIR.IR.body(mod)))
+            if MLIR.IR.dialect(op) == :llvm
+		    MLIR.API.mlirOperationDestroy(op.operation)
+		    op.operation = MLIR.API.MlirOperation(C_NULL)
+            end
+        end
+    end
 
     return Reactant.TracedUtils.CompiledMlirFnResult(
         fnwrapped,

src/TPUs.jl

@@ -38,7 +38,7 @@ function download_libtpu_if_needed(path=nothing)
     @assert path !== nothing "libtpu_dir is not set!"
     if !isfile(path * "/libtpu.so")
         Downloads.download(
-            "https://storage.googleapis.com/libtpu-nightly-releases/wheels/libtpu-nightly/libtpu_nightly-0.1.dev20250313+nightly-py3-none-manylinux_2_31_x86_64.whl",
+	    "https://storage.googleapis.com/libtpu-nightly-releases/wheels/libtpu-nightly/libtpu_nightly-0.1.dev20250326+nightly-py3-none-manylinux_2_31_x86_64.whl",
             path * "/tpu.zip",
         )
         run(`unzip -qq $(path*"/tpu.zip") -d $(path)/tmp`)