Use jit_apply in to_mlir_module

elixir-nx · Aug 13, 2024 · e6439da · e6439da
1 parent 03d7f87
commit e6439da
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/exla/lib/exla.ex b/exla/lib/exla.ex
@@ -381,12 +381,12 @@ defmodule EXLA do
       throw({:mlir_module, executable.ref})
     end
 
-    function
-    |> jit([
+    opts = [
       {EXLA, {&EXLA.Defn.LockedCache.run/2, comp_fun}},
       {:module_compilation, :to_mlir} | options
-    ])
-    |> apply(args)
+    ]
+
+    jit_apply(function, args, opts)
   catch
     {:mlir_module, ref} ->
       EXLA.MLIR.Module.as_string(%EXLA.MLIR.Module{ref: ref})