autodiff dynamic loops with conditional
```julia
using Reactant
using Enzyme
con(x) = sum(x) > 1
function f(con, x)
return ifelse.(con(x), x, x .+1)
end
function l(x, k, c)
Reactant,@trace mincut=true checkpointing=Reactant.Periodic(c) track_numbers=false for _ in 1:k
x = f(con, x)
end
return sum(x)
end
function g_l(x, dx, k, c)
Enzyme.autodiff(
Reverse, loss, Active,
Duplicated(x, dx), Const(k), Const(c)
)
return dx
end
x = Reactant.to_rarray([1.0, 2.0, 3.0])
k = ConcreteRNumber(10)
c = 3
@time closs = Reactant.@compile raise=true sync=true l(x, k, c)
@info closs(x, k, c)
dx = Enzyme.make_zero(x)
@time cgrad_loss = Reactant.@compile raise_first=true raise=true sync=true g_l(x, dx, k, c)
@info cgrad_loss(x, dx, k, c)
```
```julia
0.270359 seconds (596.73 k allocations: 33.583 MiB, 1.87% gc time, 84.15% compilation time: 15% of which was recompilation)
[ Info: ConcretePJRTNumber{Float64, 1}(6.0)
┌ Error: Compilation failed, MLIR module written to...
...
UNKNOWN: <unknown>:0: error: 'stablehlo.dynamic_pad' op can't be translated to XLA HLO
<unknown>:0: note: see current operation: %6 = "stablehlo.dynamic_pad"(%0, %2, %1, %5, %1) : (tensor<0xi1>, tensor<i1>, tensor<1xi64>, tensor<1xi64>, tensor<1xi64>) -> tensor<?xi1>
```
Ideally, one wants to be able to differentiate through dynamically-shaped loops that have conditionals inside of them, but then we generate this `dynamic_pad`. Dump from `post_xla`. Seems to be an XLA problem, but if anyone has any thoughts about workarounds for this, would be very helpful.
```
func.func @main(%arg0: tensor<3xf64> {tf.aliasing_output = 1 : i32}, %arg1: tensor<3xf64> {tf.aliasing_output = 0 : i32}, %arg2: tensor<i64>) -> (tensor<3xf64>, tensor<3xf64>) {
%c = stablehlo.constant dense<> : tensor<0xi1>
%c_0 = stablehlo.constant dense<0> : tensor<1xi64>
%c_1 = stablehlo.constant dense<false> : tensor<i1>
%c_2 = stablehlo.constant dense<0> : tensor<i64>
%cst = stablehlo.constant dense<1.000000e+00> : tensor<3xf64>
%0 = stablehlo.reshape %arg2 : (tensor<i64>) -> tensor<1xi64>
%1 = stablehlo.dynamic_pad %c, %c_1, %c_0, %0, %c_0 : (tensor<0xi1>, tensor<i1>, tensor<1xi64>, tensor<1xi64>, tensor<1xi64>) -> tensor<?xi1>
%2:3 = stablehlo.while(%iterArg = %c_2, %iterArg_3 = %arg0, %iterArg_4 = %1) : tensor<i64>, tensor<3xf64>, tensor<?xi1>
cond {
%5 = stablehlo.compare LT, %iterArg, %arg2 : (tensor<i64>, tensor<i64>) -> tensor<i1>
stablehlo.return %5 : tensor<i1>
} do {
%cst_5 = stablehlo.constant dense<1.000000e+00> : tensor<3xf64>
%cst_6 = stablehlo.constant dense<1.000000e+00> : tensor<f64>
%cst_7 = stablehlo.constant dense<0.000000e+00> : tensor<f64>
%cst_8 = stablehlo.constant dense<0.98999999999999999> : tensor<3xf64>
%c_9 = stablehlo.constant dense<1> : tensor<i64>
%5 = stablehlo.add %iterArg, %c_9 : tensor<i64>
%6 = stablehlo.multiply %iterArg_3, %cst_8 : tensor<3xf64>
%7 = stablehlo.reduce(%6 init: %cst_7) applies stablehlo.add across dimensions = [0] : (tensor<3xf64>, tensor<f64>) -> tensor<f64>
%8 = stablehlo.compare LT, %7, %cst_6 : (tensor<f64>, tensor<f64>) -> tensor<i1>
%9 = stablehlo.reshape %8 : (tensor<i1>) -> tensor<1xi1>
%10 = stablehlo.dynamic_update_slice %iterArg_4, %9, %iterArg : (tensor<?xi1>, tensor<1xi1>, tensor<i64>) -> tensor<?xi1>
%11 = stablehlo.add %6, %cst_5 : tensor<3xf64>
%12 = stablehlo.select %8, %11, %6 : tensor<i1>, tensor<3xf64>
stablehlo.return %5, %12, %10 : tensor<i64>, tensor<3xf64>, tensor<?xi1>
}
%3 = stablehlo.add %arg1, %cst : tensor<3xf64>
%4:2 = stablehlo.while(%iterArg = %c_2, %iterArg_3 = %3) : tensor<i64>, tensor<3xf64>
cond {
%5 = stablehlo.compare LT, %iterArg, %arg2 : (tensor<i64>, tensor<i64>) -> tensor<i1>
stablehlo.return %5 : tensor<i1>
} do {
%cst_4 = stablehlo.constant dense<0.98999999999999999> : tensor<3xf64>
%c_5 = stablehlo.constant dense<1> : tensor<i64>
%5 = stablehlo.add %iterArg, %c_5 : tensor<i64>
%6 = stablehlo.multiply %iterArg_3, %cst_4 : tensor<3xf64>
stablehlo.return %5, %6 : tensor<i64>, tensor<3xf64>
}
return %4#1, %2#1 : tensor<3xf64>, tensor<3xf64>
}
```
3 条评论