Compilation error when function arguments have different sizes

[NoFixedPoint]

Description

When using Reactant's @compile with Enzyme's autodiff, compilation fails if function arguments have different sizes.

Environment

• Julia Version 1.11.7
  Commit f2b3dbda30a (2025-09-08 12:10 UTC)
  Build Info:
  Official https://julialang.org/ release
  Platform Info:
  OS: Linux (x86_64-linux-gnu)
  CPU: 32 × 13th Gen Intel(R) Core(TM) i9-13900K
  WORD_SIZE: 64
  LLVM: libLLVM-16.0.6 (ORCJIT, alderlake)
  Threads: 32 default, 0 interactive, 16 GC (on 32 virtual cores)
• Reactant version: [3c362404] Reactant v0.2.169
• Enzyme version: [7da242da] Enzyme v0.13.85

Minimal Reproducible Example

using Reactant
using Enzyme

# Create arrays with different sizes
x = rand(1000)
t = rand(1001)  # Different size from x
# it works if uncommenting the following to make the size of x and t the same
# t = rand(1000)

function eval_F(x, t)
    F = x.^2  # Note: t is unused
    return F
end

function eval_jvp_x_F(v, x, t)
    jvp = Enzyme.autodiff(Forward, eval_F, Duplicated(x, v), Const(t))[1]
    return jvp
end

# This works fine
v = ones(size(x))
eval_jvp_x_F(v, x, t)

# This throws a compilation error
v_r = Reactant.to_rarray(v)
x_r = Reactant.to_rarray(x)
t_r = Reactant.to_rarray(t)
eval_jvp_x_F_compiled = @compile eval_jvp_x_F(v_r, x_r, t_r)

Expected Behavior

The compilation should succeed.

Actual Behavior

Compilation fails with an error when size(t) != size(x).

Error Message

error: type of return operand 3 ('tensor<1000xf64>') doesn't match function result type ('tensor<1001xf64>') in function @eval_jvp_x_F
┌ Error: Compilation failed, MLIR module written to /tmp/reactant_uMtOZs/module_001_4sWt_post_all_pm.mlir
└ @ Reactant.MLIR.IR ~/.julia/packages/Reactant/bR5J7/src/mlir/IR/Pass.jl:119
ERROR: "failed to run pass manager on module"
Stacktrace:
[1] run!(pm::Reactant.MLIR.IR.PassManager, mod::Reactant.MLIR.IR.Module, key::String)
@ Reactant.MLIR.IR ~/.julia/packages/Reactant/bR5J7/src/mlir/IR/Pass.jl:163
[2] run_pass_pipeline!(mod::Reactant.MLIR.IR.Module, pass_pipeline::String, key::String; enable_verifier::Bool)
@ Reactant.Compiler ~/.julia/packages/Reactant/bR5J7/src/Compiler.jl:1306
[3] run_pass_pipeline!
@ ~/.julia/packages/Reactant/bR5J7/src/Compiler.jl:1301 [inlined]
[4] compile_mlir!(mod::Reactant.MLIR.IR.Module, f::Function, args::Tuple{…}, compile_options::CompileOptions, callcache::Dict{…}, sdycache::Dict{…}; fn_kwargs::@NamedTuple{}, backend::String, runtime::Val{…}, legalize_stablehlo_to_mhlo::Bool, kwargs::@kwargs{})
@ Reactant.Compiler ~/.julia/packages/Reactant/bR5J7/src/Compiler.jl:1746
[5] compile_mlir! (repeats 2 times)
@ ~/.julia/packages/Reactant/bR5J7/src/Compiler.jl:1561 [inlined]
[6] compile_xla(f::Function, args::Tuple{…}; before_xla_optimizations::Bool, client::Nothing, serializable::Bool, kwargs::@kwargs{…})
@ Reactant.Compiler ~/.julia/packages/Reactant/bR5J7/src/Compiler.jl:3483
[7] compile_xla
@ ~/.julia/packages/Reactant/bR5J7/src/Compiler.jl:3456 [inlined]
[8] compile(f::Function, args::Tuple{…}; kwargs::@kwargs{…})
@ Reactant.Compiler ~/.julia/packages/Reactant/bR5J7/src/Compiler.jl:3555
[9] top-level scope
@ ~/.julia/packages/Reactant/bR5J7/src/Compiler.jl:2633

[NoFixedPoint]

A workaround is to define a closure and pass it as a Const when calling Enzyme.

function eval_jvp_x_F(v, x, t)
    eval_F_inline(x) = eval_F(x,t)
    jvp = Enzyme.autodiff(Forward, Const(eval_F_inline), Duplicated(x, v))[1]
    return jvp
end

[avik-pal]

Seems like tracing is fine:

module @reactant_eval_jv... attributes {mhlo.num_partitions = 1 : i64, mhlo.num_replicas = 1 : i64} {
  func.func private @"Const{typeof(eval_F)}(Main.eval_F)_autodiff"(%arg0: tensor<1000xf64> {enzymexla.memory_effects = []}, %arg1: tensor<1001xf64> {enzymexla.memory_effects = []}) -> (tensor<1000xf64>, tensor<1000xf64>, tensor<1001xf64>) attributes {enzymexla.memory_effects = []} {
    %0 = stablehlo.multiply %arg0, %arg0 : tensor<1000xf64>
    return %0, %arg0, %arg1 : tensor<1000xf64>, tensor<1000xf64>, tensor<1001xf64>
  }
  func.func @main(%arg0: tensor<1000xf64> {enzymexla.memory_effects = ["read", "write", "allocate", "free"]}, %arg1: tensor<1000xf64> {enzymexla.memory_effects = ["read", "write", "allocate", "free"], tf.aliasing_output = 1 : i32}, %arg2: tensor<1001xf64> {enzymexla.memory_effects = ["read", "write", "allocate", "free"], tf.aliasing_output = 2 : i32}) -> (tensor<1000xf64>, tensor<1000xf64>, tensor<1001xf64>) attributes {enzymexla.memory_effects = ["read", "write", "allocate", "free"]} {
    %0:3 = enzyme.fwddiff @"Const{typeof(eval_F)}(Main.eval_F)_autodiff"(%arg1, %arg0, %arg2) {activity = [#enzyme<activity enzyme_dup>, #enzyme<activity enzyme_const>], ret_activity = [#enzyme<activity enzyme_dupnoneed>, #enzyme<activity enzyme_dup>, #enzyme<activity enzyme_constnoneed>]} : (tensor<1000xf64>, tensor<1000xf64>, tensor<1001xf64>) -> (tensor<1000xf64>, tensor<1000xf64>, tensor<1001xf64>)
    return %0#0, %0#1, %0#2 : tensor<1000xf64>, tensor<1000xf64>, tensor<1001xf64>
  }
}

But after one of the enzyme passes:

// Pass pipeline:
// any(enzyme-hlo-opt{all_finite=false cse=true enable_convert_to_convolution=false max_constant_expansion=1024 max_iterations=100 no_nan=false passses=24575 top_down=false},enzyme-batch,inline{default-pipeline=canonicalize inlining-threshold=4294967295 max-iterations=4 },enzyme{postpasses= verifyPostPasses=true})
#loc1 = loc("arg1 (path=(Symbol(\22##autodiffarg#119747\22), 1))")
#loc2 = loc("arg2 (path=(Symbol(\22##autodiffarg#119747\22), 2))")
#loc4 = loc("arg1 (path=(:args, 1))")
#loc5 = loc("arg2 (path=(:args, 2))")
#loc6 = loc("arg3 (path=(:args, 3))")
"builtin.module"() ({
  "func.func"() <{function_type = (tensor<1000xf64>, tensor<1001xf64>) -> (tensor<1000xf64>, tensor<1000xf64>, tensor<1001xf64>), sym_name = "Const{typeof(eval_F)}(Main.eval_F)_autodiff", sym_visibility = "private"}> ({
  ^bb0(%arg6: tensor<1000xf64> loc("arg1 (path=(Symbol(\22##autodiffarg#119747\22), 1))"), %arg7: tensor<1001xf64> loc("arg2 (path=(Symbol(\22##autodiffarg#119747\22), 2))")):
    %5 = "stablehlo.multiply"(%arg6, %arg6) : (tensor<1000xf64>, tensor<1000xf64>) -> tensor<1000xf64> loc(#loc7)
    "func.return"(%5, %arg6, %arg7) : (tensor<1000xf64>, tensor<1000xf64>, tensor<1001xf64>) -> () loc(#loc)
  }) : () -> () loc(#loc)
  "func.func"() <{function_type = (tensor<1000xf64>, tensor<1000xf64>, tensor<1001xf64>) -> (tensor<1000xf64>, tensor<1000xf64>, tensor<1000xf64>, tensor<1001xf64>), sym_name = "eval_jvp_x_F"}> ({
  ^bb0(%arg3: tensor<1000xf64> loc("arg1 (path=(:args, 1))"), %arg4: tensor<1000xf64> loc("arg2 (path=(:args, 2))"), %arg5: tensor<1001xf64> loc("arg3 (path=(:args, 3))")):
    %4:3 = "func.call"(%arg4, %arg3, %arg5) <{callee = @"fwddiffeConst{typeof(eval_F)}(Main.eval_F)_autodiff"}> : (tensor<1000xf64>, tensor<1000xf64>, tensor<1001xf64>) -> (tensor<1000xf64>, tensor<1000xf64>, tensor<1000xf64>) loc(#loc)
    "func.return"(%4#0, %arg3, %4#1, %4#2) : (tensor<1000xf64>, tensor<1000xf64>, tensor<1000xf64>, tensor<1000xf64>) -> () loc(#loc)
  }) : () -> () loc(#loc)
  "func.func"() <{function_type = (tensor<1000xf64>, tensor<1000xf64>, tensor<1001xf64>) -> (tensor<1000xf64>, tensor<1000xf64>, tensor<1000xf64>), sym_name = "fwddiffeConst{typeof(eval_F)}(Main.eval_F)_autodiff", sym_visibility = "private"}> ({
  ^bb0(%arg0: tensor<1000xf64> loc("arg1 (path=(Symbol(\22##autodiffarg#119747\22), 1))"), %arg1: tensor<1000xf64> loc("arg1 (path=(Symbol(\22##autodiffarg#119747\22), 1))"), %arg2: tensor<1001xf64> loc("arg2 (path=(Symbol(\22##autodiffarg#119747\22), 2))")):
    %0 = "stablehlo.multiply"(%arg1, %arg0) : (tensor<1000xf64>, tensor<1000xf64>) -> tensor<1000xf64> loc(#loc7)
    %1 = "stablehlo.multiply"(%arg1, %arg0) : (tensor<1000xf64>, tensor<1000xf64>) -> tensor<1000xf64> loc(#loc7)
    %2 = "arith.addf"(%0, %1) <{fastmath = #arith.fastmath<none>}> : (tensor<1000xf64>, tensor<1000xf64>) -> tensor<1000xf64> loc(#loc7)
    %3 = "stablehlo.multiply"(%arg0, %arg0) : (tensor<1000xf64>, tensor<1000xf64>) -> tensor<1000xf64> loc(#loc7)
    "func.return"(%2, %arg0, %arg1) : (tensor<1000xf64>, tensor<1000xf64>, tensor<1000xf64>) -> () loc(#loc)
  }) : () -> () loc(#loc)
}) : () -> () loc(#loc)
#loc = loc(unknown)
#loc3 = loc("/mnt2/avik-pal/reactant/Reactant.jl/src/TracedRNumber.jl":229:0)
#loc7 = loc("*/multiply"(#loc3))