rev diff of `stablehlo.dynamic_update_slice` is broken for complex tensors

[mofeing]

it seems like #1121 broke the rev diff of stablehlo.dynamic_update_slice for complex tensors. proof of it is that the VQE benchmark (EnzymeAD/Reactant.jl#1416) worked with Reactant 0.2.138 but it doesn't on more recent versions.

this is the MLIR with optimize = false and optimize = true for Reactant 0.2.138
vqe-diff-n30-d1-v0.2.138-opt_false.mlir.txt
vqe-diff-n30-d1-v0.2.138-opt_true.mlir.txt

this is the erroring IR with optimize = false for Reactant 0.2.150
vqe-diff-n30-d1-v0.2.150-opt_false.mlir.txt

the error seems to be that the tablegen rev diff rule calls mlir::FloatType::getFloatSemantics on a complex type, which doesn't implement it?

stacktrace

[16419] signal 11 (2): Segmentation fault: 11
in expression starting at REPL[39]:1
_ZNK4mlir9FloatType17getFloatSemanticsEv at /Users/mofeing/.julia/artifacts/b67395cb97d7760a1150283970e04c23249f1230/lib/libReactantExtra.dylib (unknown line)
_ZN4mlir6enzyme15getConstantAttrENS_4TypeEN4llvm9StringRefE at /Users/mofeing/.julia/artifacts/b67395cb97d7760a1150283970e04c23249f1230/lib/libReactantExtra.dylib (unknown line)
_ZN4mlir6enzyme6detail41ReverseAutoDiffOpInterfaceInterfaceTraits13FallbackModelIN12_GLOBAL__N_133DynamicUpdateSliceOpRevDerivativeEE24createReverseModeAdjointEPKNS2_7ConceptEPNS_9OperationERNS_9OpBuilderEPNS0_21MGradientUtilsReverseEN4llvm11SmallVectorINS_5ValueELj6EEE at /Users/mofeing/.julia/artifacts/b67395cb97d7760a1150283970e04c23249f1230/lib/libReactantExtra.dylib (unknown line)
_ZN4mlir6enzyme26ReverseAutoDiffOpInterface24createReverseModeAdjointERNS_9OpBuilderEPNS0_21MGradientUtilsReverseEN4llvm11SmallVectorINS_5ValueELj6EEE at /Users/mofeing/.julia/artifacts/b67395cb97d7760a1150283970e04c23249f1230/lib/libReactantExtra.dylib (unknown line)
_ZN4mlir6enzyme12MEnzymeLogic10visitChildEPNS_9OperationERNS_9OpBuilderEPNS0_21MGradientUtilsReverseE at /Users/mofeing/.julia/artifacts/b67395cb97d7760a1150283970e04c23249f1230/lib/libReactantExtra.dylib (unknown line)
_ZN4mlir6enzyme12MEnzymeLogic13differentiateEPNS0_21MGradientUtilsReverseERNS_6RegionES5_N4llvm12function_refIFvRNS_9OpBuilderEPNS_5BlockEEEENSt3__18functionIFNSE_4pairINS_5ValueESH_EENS_4TypeEEEE at /Users/mofeing/.julia/artifacts/b67395cb97d7760a1150283970e04c23249f1230/lib/libReactantExtra.dylib (unknown line)
_ZN4mlir6enzyme12MEnzymeLogic17CreateReverseDiffENS_19FunctionOpInterfaceENSt3__16vectorI10DIFFE_TYPENS3_9allocatorIS5_EEEES8_RNS0_13MTypeAnalysisENS4_IbNS6_IbEEEESC_14DerivativeModebmNS_4TypeENS0_11MFnTypeInfoESC_PvbN4llvm9StringRefEbb at /Users/mofeing/.julia/artifacts/b67395cb97d7760a1150283970e04c23249f1230/lib/libReactantExtra.dylib (unknown line)
_ZN12_GLOBAL__N_117DifferentiatePass16lowerEnzymeCallsERN4mlir21SymbolTableCollectionENS1_19FunctionOpInterfaceE at /Users/mofeing/.julia/artifacts/b67395cb97d7760a1150283970e04c23249f1230/lib/libReactantExtra.dylib (unknown line)
_ZN4mlir6detail4walkINS_15ForwardIteratorEEEvPNS_9OperationEN4llvm12function_refIFvS4_EEENS_9WalkOrderE at /Users/mofeing/.julia/artifacts/b67395cb97d7760a1150283970e04c23249f1230/lib/libReactantExtra.dylib (unknown line)
_ZN12_GLOBAL__N_117DifferentiatePass14runOnOperationEv at /Users/mofeing/.julia/artifacts/b67395cb97d7760a1150283970e04c23249f1230/lib/libReactantExtra.dylib (unknown line)
_ZN4mlir6detail17OpToOpPassAdaptor3runEPNS_4PassEPNS_9OperationENS_15AnalysisManagerEbj at /Users/mofeing/.julia/artifacts/b67395cb97d7760a1150283970e04c23249f1230/lib/libReactantExtra.dylib (unknown line)
_ZN4mlir11PassManager9runPassesEPNS_9OperationENS_15AnalysisManagerE at /Users/mofeing/.julia/artifacts/b67395cb97d7760a1150283970e04c23249f1230/lib/libReactantExtra.dylib (unknown line)
_ZN4mlir11PassManager3runEPNS_9OperationE at /Users/mofeing/.julia/artifacts/b67395cb97d7760a1150283970e04c23249f1230/lib/libReactantExtra.dylib (unknown line)
mlirPassManagerRunOnOp at /Users/mofeing/.julia/artifacts/b67395cb97d7760a1150283970e04c23249f1230/lib/libReactantExtra.dylib (unknown line)
jlplt_mlirPassManagerRunOnOp_28272 at /Users/mofeing/.julia/compiled/v1.11/Reactant/p9PzF_rBHx1.dylib (unknown line)
mlirPassManagerRunOnOp at /Users/mofeing/Developer/Reactant.jl/src/mlir/libMLIR_h.jl:8876 [inlined]
run! at /Users/mofeing/Developer/Reactant.jl/src/mlir/IR/Pass.jl:154
#run_pass_pipeline!#2 at /Users/mofeing/Developer/Reactant.jl/src/Compiler.jl:1240

cc @wsmoses

[mofeing]

I managed to get a MWE

using Reactant
using Reactant.TracedUtils: broadcast_to_size
using Enzyme
Reactant.MLIR.IR.DUMP_MLIR_ALWAYS[] = true

function f(A, θ)
    M = copy(A)
    @allowscalar begin
        c = cos(θ / 2)
        s = -im * sin(θ / 2)
        M[1, 1] = c
        M[2, 2] = s
    end
    return sum(M * M')
end

Are = Reactant.ConcreteRArray(ones(ComplexF64, 2, 2))
θre = Reactant.ConcreteRNumber(π / 2)

@code_hlo optimize = false f(Are, θre)
@jit f(Are, θre)

function ∇f(A, θ)
    return Enzyme.gradient(ReverseWithPrimal, f, A, θ)
end

Reactant.@skip_rewrite_func ∇f

@code_hlo optimize = false ∇f(Are, θre)

compile_options = Reactant.Compiler.CompileOptions()
opt_passes = Reactant.Compiler.optimization_passes(
    compile_options; sroa=true, recognize_comms=false, lower_comms=false
)
opt_passes2 = Reactant.Compiler.optimization_passes(
    compile_options; sroa=false, recognize_comms=false, lower_comms=false
)

optimize = join([
    opt_passes,
    "enzyme-batch",
    opt_passes2,
    "enzyme{}",
], ',')

@code_hlo optimize = optimize ∇f(Are, θre)

The custom optimize list tries to reproduce the passes in :before_enzyme.

@wsmoses Here I'm leaving the MLIR of join([opt_passes, "enzyme-batch", opt_passes2]). You can reproduce the segmentation fault by just adding the last enzyme{} pass.

module @"reactant_\E2\88\87f" attributes {mhlo.num_partitions = 1 : i64, mhlo.num_replicas = 1 : i64} {
  func.func private @"Const{typeof(f)}(Main.f)_autodiff"(%arg0: tensor<2x2xcomplex<f64>>, %arg1: tensor<f64>) -> (tensor<complex<f64>>, tensor<2x2xcomplex<f64>>, tensor<f64>) {
    %cst = stablehlo.constant dense<5.000000e-01> : tensor<1x1xf64>
    %cst_0 = stablehlo.constant dense<(0.000000e+00,-1.000000e+00)> : tensor<1x1xcomplex<f64>>
    %cst_1 = stablehlo.constant dense<(0.000000e+00,0.000000e+00)> : tensor<complex<f64>>
    %c = stablehlo.constant dense<1> : tensor<i32>
    %c_2 = stablehlo.constant dense<0> : tensor<i32>
    %0 = stablehlo.transpose %arg0, dims = [1, 0] : (tensor<2x2xcomplex<f64>>) -> tensor<2x2xcomplex<f64>>
    %1 = stablehlo.reshape %arg1 : (tensor<f64>) -> tensor<1x1xf64>
    %2 = stablehlo.multiply %1, %cst : tensor<1x1xf64>
    %3 = stablehlo.cosine %2 : tensor<1x1xf64>
    %4 = stablehlo.convert %3 : (tensor<1x1xf64>) -> tensor<1x1xcomplex<f64>>
    %5 = stablehlo.dynamic_update_slice %0, %4, %c_2, %c_2 : (tensor<2x2xcomplex<f64>>, tensor<1x1xcomplex<f64>>, tensor<i32>, tensor<i32>) -> tensor<2x2xcomplex<f64>>
    %6 = stablehlo.sine %2 : tensor<1x1xf64>
    %7 = stablehlo.convert %6 : (tensor<1x1xf64>) -> tensor<1x1xcomplex<f64>>
    %8 = stablehlo.multiply %cst_0, %7 : tensor<1x1xcomplex<f64>>
    %9 = stablehlo.dynamic_update_slice %5, %8, %c, %c : (tensor<2x2xcomplex<f64>>, tensor<1x1xcomplex<f64>>, tensor<i32>, tensor<i32>) -> tensor<2x2xcomplex<f64>>
    %10 = stablehlo.transpose %9, dims = [1, 0] : (tensor<2x2xcomplex<f64>>) -> tensor<2x2xcomplex<f64>>
    %11 = chlo.conj %10 : tensor<2x2xcomplex<f64>> -> tensor<2x2xcomplex<f64>>
    %12 = stablehlo.dot_general %9, %11, contracting_dims = [1] x [0], precision = [DEFAULT, DEFAULT] : (tensor<2x2xcomplex<f64>>, tensor<2x2xcomplex<f64>>) -> tensor<2x2xcomplex<f64>>
    %13 = stablehlo.reduce(%12 init: %cst_1) applies stablehlo.add across dimensions = [0, 1] : (tensor<2x2xcomplex<f64>>, tensor<complex<f64>>) -> tensor<complex<f64>>
    return %13, %arg0, %arg1 : tensor<complex<f64>>, tensor<2x2xcomplex<f64>>, tensor<f64>
  }
  func.func @main(%arg0: tensor<2x2xcomplex<f64>> {tf.aliasing_output = 3 : i32}, %arg1: tensor<f64> {tf.aliasing_output = 4 : i32}) -> (tensor<2x2xcomplex<f64>>, tensor<f64>, tensor<complex<f64>>, tensor<2x2xcomplex<f64>>, tensor<f64>) {
    %cst = stablehlo.constant dense<(1.000000e+00,0.000000e+00)> : tensor<complex<f64>>
    %cst_0 = stablehlo.constant dense<0.000000e+00> : tensor<f64>
    %cst_1 = stablehlo.constant dense<(0.000000e+00,0.000000e+00)> : tensor<2x2xcomplex<f64>>
    %0:5 = enzyme.autodiff @"Const{typeof(f)}(Main.f)_autodiff"(%arg0, %arg1, %cst, %cst_1, %cst_0) {activity = [#enzyme<activity enzyme_active>, #enzyme<activity enzyme_active>], ret_activity = [#enzyme<activity enzyme_active>, #enzyme<activity enzyme_active>, #enzyme<activity enzyme_active>]} : (tensor<2x2xcomplex<f64>>, tensor<f64>, tensor<complex<f64>>, tensor<2x2xcomplex<f64>>, tensor<f64>) -> (tensor<complex<f64>>, tensor<2x2xcomplex<f64>>, tensor<f64>, tensor<2x2xcomplex<f64>>, tensor<f64>)
    return %0#3, %0#4, %0#0, %0#1, %0#2 : tensor<2x2xcomplex<f64>>, tensor<f64>, tensor<complex<f64>>, tensor<2x2xcomplex<f64>>, tensor<f64>
  }
}

pipeline without failing pass:

inline{default-pipeline=canonicalize max-iterations=4},propagate-constant-bounds,sroa-wrappers{instcombine=false instsimplify=true },canonicalize,sroa-wrappers{instcombine=false instsimplify=true },libdevice-funcs-raise,canonicalize,remove-duplicate-func-def,canonicalize,cse,canonicalize,enzyme-hlo-generate-td{patterns=compare_op_canon<16>;transpose_transpose<16>;broadcast_in_dim_op_canon<16>;convert_op_canon<16>;dynamic_broadcast_in_dim_op_not_actually_dynamic<16>;chained_dynamic_broadcast_in_dim_canonicalization<16>;dynamic_broadcast_in_dim_all_dims_non_expanding<16>;noop_reduce_op_canon<16>;empty_reduce_op_canon<16>;dynamic_reshape_op_canon<16>;get_tuple_element_op_canon<16>;real_op_canon<16>;imag_op_canon<16>;conj_complex_negate<16>;get_dimension_size_op_canon<16>;reshape_op_canon<16>;merge_consecutive_reshapes<16>;transpose_is_reshape<16>;zero_extent_tensor_canon<16>;cse_broadcast_in_dim<16>;cse_slice<16>;cse_transpose<16>;cse_convert<16>;cse_dot_general<16>;cse_reshape<16>;cse_mul<16>;cse_div<16>;cse_add<16>;cse_subtract<16>;cse_min<16>;cse_max<16>;cse_neg<16>;cse_abs<16>;cse_concatenate<16>;concatenate_op_canon<16>(1024);select_op_canon<16>(1024);add_simplify<16>;sub_simplify<16>;and_simplify<16>;max_simplify<16>;min_simplify<16>;or_simplify<16>;xor_simplify<16>;mul_simplify<16>;div_simplify<16>;rem_simplify<16>;pow_simplify<16>;simplify_extend<16>;simplify_wrap<16>;simplify_rotate<16>;noop_slice<16>;noop_reverse<16>;slice_slice<16>;shift_right_logical_simplify<16>;slice_simplify<16>;convert_simplify<16>;dynamic_slice_to_static<16>;dynamic_update_slice_elim<16>;concat_to_broadcast<16>;reduce_to_reshape<16>;broadcast_to_reshape<16>;slice_internal;iota_simplify<16>(1024);broadcast_in_dim_simplify<16>(1024);convert_concat<1>;dynamic_update_to_concat<1>;slice_of_dynamic_update<1>;slice_elementwise<1>;dot_reshape_dot<1>;concat_fuse<1>;concat_push_binop_add<1>;concat_push_binop_mul<1>;reduce_concat<1>;slice_concat<1>;concat_slice<1>;select_op_used_within_if<1>;bin_broadcast_splat_add<1>;bin_broadcast_splat_subtract<1>;bin_broadcast_splat_div<1>;bin_broadcast_splat_mul<1>;dot_general_simplify<16>;transpose_simplify<16>;reshape_empty_broadcast<1>;broadcast_reshape<1>;transpose_dot_reorder<1>;dot_transpose<1>;transpose_convolution<1>;convolution_transpose<1>;convert_convert_float<1>;reshape_iota<1>;broadcast_reduce<1>;slice_dot_general<1>;if_inline<1>;if_to_select<1>;divide_sqrt_to_multiply_rsqrt<16>;associative_binary_op_reordering<1>;transpose_broadcast_in_dim_to_broadcast_in_dim<16>;replace_neg_add_with_subtract;binop_const_simplify;not_select_simplify;common_compare_expression_rewrite;compare_select_simplify;while_simplify<1>(1);if_remove_unused;transpose_reshape_to_broadcast;reshape_transpose_to_broadcast;dus_dus;dus_dus_concat;abs_positive_simplify;transpose_unary_transpose_abs;transpose_unary_transpose_neg;transpose_unary_transpose_sqrt;transpose_unary_transpose_rsqrt;transpose_unary_transpose_ceil;transpose_unary_transpose_convert;transpose_unary_transpose_cosine;transpose_unary_transpose_exp;transpose_unary_transpose_expm1;transpose_unary_transpose_log;transpose_unary_transpose_log1p;transpose_unary_transpose_sign;transpose_unary_transpose_sine;transpose_unary_transpose_tanh;select_comp_iota_const_simplify<1>;sign_abs_simplify<1>;broadcastindim_is_reshape;slice_reduce_window<1>;while_deadresult;while_dus;dus_licm(0);while_op_induction_replacement;dus_concat;slice_dus_to_concat;while_induction_reduction;slice_licm(0);elementwise_licm(0);concatenate_licm(0);slice_broadcast;while_licm<1>(1);associative_common_mul_op_reordering;slice_select_to_select_slice;slice_if;dus_to_i32;slice_extend;concat_wrap;cse_extend<16>;cse_wrap<16>;cse_rotate<16>;cse_rotate<16>;concat_concat_axis_swap;concat_concat_to_dus;broadcast_iota_simplify;select_comp_iota_to_dus;compare_cleanup;broadcast_compare;not_compare;broadcast_iota;cse_iota;compare_iota_const_simplify;reshuffle_ands_compares;square_abs_simplify;divide_divide_simplify;concat_reshape_slice;full_reduce_reshape_or_transpose;concat_reshape_reduce;concat_elementwise;reduce_reduce;conj_real;select_broadcast_in_dim;if_op_lift_common_ops;involution_neg_simplify;involution_conj_simplify;involution_not_simplify;real_conj_simplify;conj_complex_simplify;split_convolution_into_reverse_convolution;power_multiply_to_power;log_simplify;neg_mul_const_simplify;neg_div_const_simplify;reshape_deletions_broadcast_in_dim_simplify;reshape_insertions_broadcast_in_dim_simplify;scatter_to_dynamic_update_slice<1>;scatter_multiply_simplify;unary_elementwise_scatter_simplify;scatter_indices_are_unique;scatter_update_computation_const_prop;dynamic_gather_op_is_not_dynamic<16>;gather_op_canon<16>;gather_elementwise;gather_const_prop;dus_pad;cse_pad<16>;pad_simplify<16>(1024);select_pad_to_dus<1>;and_pad_pad<1>;negative_pad_to_slice<16>;slice_pad<1>;pad_reshape_pad<1>;pad_pad<1>;add_pad_pad_to_concat<1>;concat_pad<1>;reduce_pad<1>;broadcast_pad<1>;zero_product_reshape_pad<1>;mul_zero_pad<1>;div_zero_pad<1>;binop_const_reshape_pad<1>;binop_const_pad_add<1>;binop_const_pad_subtract<1>;binop_const_pad_mul<1>;binop_const_pad_div<1>;binop_binop_pad_pad_add<1>;binop_binop_pad_pad_mul<1>;binop_pad_pad_add<1>;binop_pad_pad_subtract<1>;binop_pad_pad_mul<1>;binop_pad_pad_div<1>;binop_pad_pad_min<1>;binop_pad_pad_max<1>;unary_pad_push_convert<1>;unary_pad_push_tanh<1>;unary_pad_push_exp<1>;concat_to_pad<1>;pad_licm(0);while_pad_induction_reduction;pad_concat_to_concat_pad;rotate_pad;concat_multipad;speculate_if_pad_to_select;chlo_inf_const_prop<16>;gamma_const_prop<16>;abs_const_prop<16>;log_const_prop<1>;log_plus_one_const_prop<1>;is_finite_const_prop;not_const_prop;neg_const_prop;sqrt_const_prop;rsqrt_const_prop;cos_const_prop;sin_const_prop;exp_const_prop;expm1_const_prop;tanh_const_prop;logistic_const_prop;conj_const_prop;ceil_const_prop;cbrt_const_prop;real_const_prop;imag_const_prop;round_const_prop;round_nearest_even_const_prop;sign_const_prop;floor_const_prop;tan_const_prop;add_const_prop;and_const_prop;atan2_const_prop;complex_const_prop;div_const_prop;max_const_prop;min_const_prop;mul_const_prop;or_const_prop;pow_const_prop;rem_const_prop;sub_const_prop;xor_const_prop;const_prop_through_barrier<16>;concat_const_prop<1>(1024);dynamic_update_slice_const_prop(1024);dus_slice_simplify;reshape_concat;reshape_dus;dot_reshape_pad<1>;pad_dot_general<1>(0);pad_dot_general<1>(1);reshape_pad;reshape_wrap;reshape_rotate;reshape_extend;reshape_slice(1);reshape_elementwise(1);transpose_select;transpose_while;transpose_slice;transpose_concat;transpose_iota;transpose_reduce;transpose_reduce_window;transpose_dus;transpose_pad<1>;transpose_einsum<1>;transpose_wrap;transpose_extend;transpose_rotate;transpose_dynamic_slice;transpose_reverse;transpose_batch_norm_training;transpose_batch_norm_inference;transpose_batch_norm_grad;transpose_if;transpose_fft;transpose_reshape;transpose_elementwise(1);no_nan_add_sub_simplify(0);no_nan_mul_simplify(0);no_nan_div_simplify(0)},transform-interpreter,enzyme-hlo-remove-transform,enzyme-batch,inline{default-pipeline=canonicalize max-iterations=4},canonicalize,cse,canonicalize,enzyme-hlo-generate-td{patterns=compare_op_canon<16>;transpose_transpose<16>;broadcast_in_dim_op_canon<16>;convert_op_canon<16>;dynamic_broadcast_in_dim_op_not_actually_dynamic<16>;chained_dynamic_broadcast_in_dim_canonicalization<16>;dynamic_broadcast_in_dim_all_dims_non_expanding<16>;noop_reduce_op_canon<16>;empty_reduce_op_canon<16>;dynamic_reshape_op_canon<16>;get_tuple_element_op_canon<16>;real_op_canon<16>;imag_op_canon<16>;conj_complex_negate<16>;get_dimension_size_op_canon<16>;reshape_op_canon<16>;merge_consecutive_reshapes<16>;transpose_is_reshape<16>;zero_extent_tensor_canon<16>;cse_broadcast_in_dim<16>;cse_slice<16>;cse_transpose<16>;cse_convert<16>;cse_dot_general<16>;cse_reshape<16>;cse_mul<16>;cse_div<16>;cse_add<16>;cse_subtract<16>;cse_min<16>;cse_max<16>;cse_neg<16>;cse_abs<16>;cse_concatenate<16>;concatenate_op_canon<16>(1024);select_op_canon<16>(1024);add_simplify<16>;sub_simplify<16>;and_simplify<16>;max_simplify<16>;min_simplify<16>;or_simplify<16>;xor_simplify<16>;mul_simplify<16>;div_simplify<16>;rem_simplify<16>;pow_simplify<16>;simplify_extend<16>;simplify_wrap<16>;simplify_rotate<16>;noop_slice<16>;noop_reverse<16>;slice_slice<16>;shift_right_logical_simplify<16>;slice_simplify<16>;convert_simplify<16>;dynamic_slice_to_static<16>;dynamic_update_slice_elim<16>;concat_to_broadcast<16>;reduce_to_reshape<16>;broadcast_to_reshape<16>;slice_internal;iota_simplify<16>(1024);broadcast_in_dim_simplify<16>(1024);convert_concat<1>;dynamic_update_to_concat<1>;slice_of_dynamic_update<1>;slice_elementwise<1>;dot_reshape_dot<1>;concat_fuse<1>;concat_push_binop_add<1>;concat_push_binop_mul<1>;reduce_concat<1>;slice_concat<1>;concat_slice<1>;select_op_used_within_if<1>;bin_broadcast_splat_add<1>;bin_broadcast_splat_subtract<1>;bin_broadcast_splat_div<1>;bin_broadcast_splat_mul<1>;dot_general_simplify<16>;transpose_simplify<16>;reshape_empty_broadcast<1>;broadcast_reshape<1>;transpose_dot_reorder<1>;dot_transpose<1>;transpose_convolution<1>;convolution_transpose<1>;convert_convert_float<1>;reshape_iota<1>;broadcast_reduce<1>;slice_dot_general<1>;if_inline<1>;if_to_select<1>;divide_sqrt_to_multiply_rsqrt<16>;associative_binary_op_reordering<1>;transpose_broadcast_in_dim_to_broadcast_in_dim<16>;replace_neg_add_with_subtract;binop_const_simplify;not_select_simplify;common_compare_expression_rewrite;compare_select_simplify;while_simplify<1>(1);if_remove_unused;transpose_reshape_to_broadcast;reshape_transpose_to_broadcast;dus_dus;dus_dus_concat;abs_positive_simplify;transpose_unary_transpose_abs;transpose_unary_transpose_neg;transpose_unary_transpose_sqrt;transpose_unary_transpose_rsqrt;transpose_unary_transpose_ceil;transpose_unary_transpose_convert;transpose_unary_transpose_cosine;transpose_unary_transpose_exp;transpose_unary_transpose_expm1;transpose_unary_transpose_log;transpose_unary_transpose_log1p;transpose_unary_transpose_sign;transpose_unary_transpose_sine;transpose_unary_transpose_tanh;select_comp_iota_const_simplify<1>;sign_abs_simplify<1>;broadcastindim_is_reshape;slice_reduce_window<1>;while_deadresult;while_dus;dus_licm(0);while_op_induction_replacement;dus_concat;slice_dus_to_concat;while_induction_reduction;slice_licm(0);elementwise_licm(0);concatenate_licm(0);slice_broadcast;while_licm<1>(1);associative_common_mul_op_reordering;slice_select_to_select_slice;slice_if;dus_to_i32;slice_extend;concat_wrap;cse_extend<16>;cse_wrap<16>;cse_rotate<16>;cse_rotate<16>;concat_concat_axis_swap;concat_concat_to_dus;broadcast_iota_simplify;select_comp_iota_to_dus;compare_cleanup;broadcast_compare;not_compare;broadcast_iota;cse_iota;compare_iota_const_simplify;reshuffle_ands_compares;square_abs_simplify;divide_divide_simplify;concat_reshape_slice;full_reduce_reshape_or_transpose;concat_reshape_reduce;concat_elementwise;reduce_reduce;conj_real;select_broadcast_in_dim;if_op_lift_common_ops;involution_neg_simplify;involution_conj_simplify;involution_not_simplify;real_conj_simplify;conj_complex_simplify;split_convolution_into_reverse_convolution;power_multiply_to_power;log_simplify;neg_mul_const_simplify;neg_div_const_simplify;reshape_deletions_broadcast_in_dim_simplify;reshape_insertions_broadcast_in_dim_simplify;scatter_to_dynamic_update_slice<1>;scatter_multiply_simplify;unary_elementwise_scatter_simplify;scatter_indices_are_unique;scatter_update_computation_const_prop;dynamic_gather_op_is_not_dynamic<16>;gather_op_canon<16>;gather_elementwise;gather_const_prop;dus_pad;cse_pad<16>;pad_simplify<16>(1024);select_pad_to_dus<1>;and_pad_pad<1>;negative_pad_to_slice<16>;slice_pad<1>;pad_reshape_pad<1>;pad_pad<1>;add_pad_pad_to_concat<1>;concat_pad<1>;reduce_pad<1>;broadcast_pad<1>;zero_product_reshape_pad<1>;mul_zero_pad<1>;div_zero_pad<1>;binop_const_reshape_pad<1>;binop_const_pad_add<1>;binop_const_pad_subtract<1>;binop_const_pad_mul<1>;binop_const_pad_div<1>;binop_binop_pad_pad_add<1>;binop_binop_pad_pad_mul<1>;binop_pad_pad_add<1>;binop_pad_pad_subtract<1>;binop_pad_pad_mul<1>;binop_pad_pad_div<1>;binop_pad_pad_min<1>;binop_pad_pad_max<1>;unary_pad_push_convert<1>;unary_pad_push_tanh<1>;unary_pad_push_exp<1>;concat_to_pad<1>;pad_licm(0);while_pad_induction_reduction;pad_concat_to_concat_pad;rotate_pad;concat_multipad;speculate_if_pad_to_select;chlo_inf_const_prop<16>;gamma_const_prop<16>;abs_const_prop<16>;log_const_prop<1>;log_plus_one_const_prop<1>;is_finite_const_prop;not_const_prop;neg_const_prop;sqrt_const_prop;rsqrt_const_prop;cos_const_prop;sin_const_prop;exp_const_prop;expm1_const_prop;tanh_const_prop;logistic_const_prop;conj_const_prop;ceil_const_prop;cbrt_const_prop;real_const_prop;imag_const_prop;round_const_prop;round_nearest_even_const_prop;sign_const_prop;floor_const_prop;tan_const_prop;add_const_prop;and_const_prop;atan2_const_prop;complex_const_prop;div_const_prop;max_const_prop;min_const_prop;mul_const_prop;or_const_prop;pow_const_prop;rem_const_prop;sub_const_prop;xor_const_prop;const_prop_through_barrier<16>;concat_const_prop<1>(1024);dynamic_update_slice_const_prop(1024);dus_slice_simplify;reshape_concat;reshape_dus;dot_reshape_pad<1>;pad_dot_general<1>(0);pad_dot_general<1>(1);reshape_pad;reshape_wrap;reshape_rotate;reshape_extend;reshape_slice(1);reshape_elementwise(1);transpose_select;transpose_while;transpose_slice;transpose_concat;transpose_iota;transpose_reduce;transpose_reduce_window;transpose_dus;transpose_pad<1>;transpose_einsum<1>;transpose_wrap;transpose_extend;transpose_rotate;transpose_dynamic_slice;transpose_reverse;transpose_batch_norm_training;transpose_batch_norm_inference;transpose_batch_norm_grad;transpose_if;transpose_fft;transpose_reshape;transpose_elementwise(1);no_nan_add_sub_simplify(0);no_nan_mul_simplify(0);no_nan_div_simplify(0)},transform-interpreter,enzyme-hlo-remove-transform