Weird indexing error with similar + fill! #1602
Description
Hello, again related to https://discourse.julialang.org/t/using-reactant-with-lux-and-enzyme-to-speed-up-training-in-physics-context/131898/17
mwe :
using Reactant
function foo(X)
A = similar(X,10)
fill!(A, 2.0)
A[2:4] .+= 2
return A
end
X = Reactant.to_rarray(rand(4,5))
println(@jit foo(X))error
ERROR: LoadError: BoundsError: attempt to access 10-element Reactant.TracedRArray{Float64, 1} at index [2, 3, 4]
Stacktrace:
[1] throw_boundserror(A::Reactant.TracedRArray{Float64, 1}, I::Tuple{Int64, Int64, Int64})
@ Base ./essentials.jl:14
[2] checkbounds
@ ./abstractarray.jl:699 [inlined]
[3] _setindex!
@ ./abstractarray.jl:1435 [inlined]
[4] setindex!
@ ./abstractarray.jl:1413 [inlined]
[5] (::Nothing)(none::typeof(setindex!), none::Reactant.TracedRArray{Float64, 1}, none::Reactant.TracedRArray{Float64, 1}, none::Tuple{CartesianIndex{1}, CartesianIndex{1}, CartesianIndex{1}})
@ Reactant ./<missing>:0
[6] getindex
@ ./tuple.jl:31 [inlined]
[7] to_indices
@ ./multidimensional.jl:860 [inlined]
[8] to_indices
@ ./multidimensional.jl:855 [inlined]
[9] setindex!
@ ./abstractarray.jl:1413 [inlined]
[10] call_with_reactant(::typeof(setindex!), ::Reactant.TracedRArray{Float64, 1}, ::Reactant.TracedRArray{Float64, 1}, ::CartesianIndex{1}, ::CartesianIndex{1}, ::CartesianIndex{1})
@ Reactant ~/.julia/packages/Reactant/gBXlB/src/utils.jl:0
[11] macro expansion
@ ~/.julia/packages/Reactant/gBXlB/src/utils.jl:293 [inlined]
[12] applyiterate_with_reactant(::typeof(iterate), ::typeof(setindex!), ::Tuple{Reactant.TracedRArray{Float64, 1}, Reactant.TracedRArray{Float64, 1}}, ::Vector{CartesianIndex{1}})
@ Reactant ~/.julia/packages/Reactant/gBXlB/src/utils.jl:279
[13] set_mlir_data!
@ ~/.julia/packages/Reactant/gBXlB/src/TracedUtils.jl:153 [inlined]
[14] _copyto!
@ ~/.julia/packages/Reactant/gBXlB/src/TracedRArray.jl:775 [inlined]
[15] (::Nothing)(none::typeof(Reactant.TracedRArrayOverrides._copyto!), none::SubArray{Reactant.TracedRNumber{Float64}, 1, Reactant.TracedRArray{Float64, 1}, Tuple{UnitRange{Int64}}, true}, none::Base.Broadcast.Broadcasted{Reactant.TracedRArrayOverrides.AbstractReactantArrayStyle{1}, Tuple{Base.OneTo{Int64}}, typeof(+), Tuple{Reactant.TracedRArray{Float64, 1}, Int64}})
@ Reactant ./<missing>:0
[16] getproperty
@ ./Base.jl:49 [inlined]
[17] axes
@ ./subarray.jl:504 [inlined]
[18] _copyto!
@ ~/.julia/packages/Reactant/gBXlB/src/TracedRArray.jl:764 [inlined]
[19] call_with_reactant(::typeof(Reactant.TracedRArrayOverrides._copyto!), ::SubArray{Reactant.TracedRNumber{Float64}, 1, Reactant.TracedRArray{Float64, 1}, Tuple{UnitRange{Int64}}, true}, ::Base.Broadcast.Broadcasted{Reactant.TracedRArrayOverrides.AbstractReactantArrayStyle{1}, Tuple{Base.OneTo{Int64}}, typeof(+), Tuple{Reactant.TracedRArray{Float64, 1}, Int64}})
@ Reactant ~/.julia/packages/Reactant/gBXlB/src/utils.jl:0
[20] materialize!
@ ~/.julia/packages/Reactant/gBXlB/src/TracedRArray.jl:722 [inlined]
[21] materialize!
@ ./broadcast.jl:880 [inlined]
[22] (::Nothing)(none::typeof(Base.Broadcast.materialize!), none::SubArray{Reactant.TracedRNumber{Float64}, 1, Reactant.TracedRArray{Float64, 1}, Tuple{UnitRange{Int64}}, true}, none::Base.Broadcast.Broadcasted{Reactant.TracedRArrayOverrides.AbstractReactantArrayStyle{1}, Nothing, typeof(+), Tuple{Reactant.TracedRArray{Float64, 1}, Int64}})
@ Reactant ./<missing>:0
[23] getproperty
@ ./Base.jl:49 [inlined]
[24] materialize!
@ ~/.julia/packages/Reactant/gBXlB/src/TracedRArray.jl:722 [inlined]
[25] materialize!
@ ./broadcast.jl:880 [inlined]
[26] call_with_reactant(::typeof(Base.Broadcast.materialize!), ::SubArray{Reactant.TracedRNumber{Float64}, 1, Reactant.TracedRArray{Float64, 1}, Tuple{UnitRange{Int64}}, true}, ::Base.Broadcast.Broadcasted{Reactant.TracedRArrayOverrides.AbstractReactantArrayStyle{1}, Nothing, typeof(+), Tuple{Reactant.TracedRArray{Float64, 1}, Int64}})
@ Reactant ~/.julia/packages/Reactant/gBXlB/src/utils.jl:0
[27] foo
@ /mnt/c/Users/yolha/Desktop/juju_tests/mini/test/main2.jl:6 [inlined]
[28] (::Nothing)(none::typeof(foo), none::Reactant.TracedRArray{Float64, 2})
@ Reactant ./<missing>:0
[29] similar
@ ./abstractarray.jl:823 [inlined]
[30] foo
@ /mnt/c/Users/yolha/Desktop/juju_tests/mini/test/main2.jl:4 [inlined]
[31] call_with_reactant(::typeof(foo), ::Reactant.TracedRArray{Float64, 2})
@ Reactant ~/.julia/packages/Reactant/gBXlB/src/utils.jl:0
[32] make_mlir_fn(f::typeof(foo), args::Tuple{ConcretePJRTArray{Float64, 2, 1, Reactant.Sharding.ShardInfo{Reactant.Sharding.NoSharding, Nothing}}}, kwargs::@NamedTuple{}, name::String, concretein::Bool; toscalar::Bool, return_dialect::Symbol, args_in_result::Symbol, construct_function_without_args::Bool, do_transpose::Bool, input_shardings::Nothing, output_shardings::Nothing, runtime::Val{:PJRT}, verify_arg_names::Nothing, argprefix::Symbol, resprefix::Symbol, resargprefix::Symbol, num_replicas::Int64, optimize_then_pad::Bool)
@ Reactant.TracedUtils ~/.julia/packages/Reactant/gBXlB/src/TracedUtils.jl:332
[33] compile_mlir!(mod::Reactant.MLIR.IR.Module, f::Function, args::Tuple{ConcretePJRTArray{Float64, 2, 1, Reactant.Sharding.ShardInfo{Reactant.Sharding.NoSharding, Nothing}}}, compile_options::CompileOptions, callcache::Dict{Vector, @NamedTuple{f_name::String, mlir_result_types::Vector{Reactant.MLIR.IR.Type}, traced_result, mutated_args::Vector{Int64}, linear_results::Vector{Union{ReactantCore.MissingTracedValue, Reactant.TracedRArray, Reactant.TracedRNumber}}, fnwrapped::Bool, argprefix::Symbol, resprefix::Symbol, resargprefix::Symbol}}, sdycache::Dict{Tuple{AbstractVector{Int64}, NTuple{var"#s1734", Symbol} where var"#s1734", NTuple{N, Int64} where N}, @NamedTuple{sym_name::Reactant.MLIR.IR.Attribute, mesh_attr::Reactant.MLIR.IR.Attribute, mesh_op::Reactant.MLIR.IR.Operation, mesh::Reactant.Sharding.Mesh}}; fn_kwargs::@NamedTuple{}, backend::String, runtime::Val{:PJRT}, legalize_stablehlo_to_mhlo::Bool, kwargs::@Kwargs{})
@ Reactant.Compiler ~/.julia/packages/Reactant/gBXlB/src/Compiler.jl:1555
[34] compile_mlir! (repeats 2 times)
@ ~/.julia/packages/Reactant/gBXlB/src/Compiler.jl:1522 [inlined]
[35] compile_xla(f::Function, args::Tuple{ConcretePJRTArray{Float64, 2, 1, Reactant.Sharding.ShardInfo{Reactant.Sharding.NoSharding, Nothing}}}; before_xla_optimizations::Bool, client::Nothing, serializable::Bool, kwargs::@Kwargs{compile_options::CompileOptions, fn_kwargs::@NamedTuple{}})
@ Reactant.Compiler ~/.julia/packages/Reactant/gBXlB/src/Compiler.jl:3433
[36] compile_xla
@ ~/.julia/packages/Reactant/gBXlB/src/Compiler.jl:3406 [inlined]
[37] compile(f::Function, args::Tuple{ConcretePJRTArray{Float64, 2, 1, Reactant.Sharding.ShardInfo{Reactant.Sharding.NoSharding, Nothing}}}; kwargs::@Kwargs{fn_kwargs::@NamedTuple{}, client::Nothing, reshape_propagate::Symbol, raise_first::Bool, assert_nonallocating::Bool, legalize_chlo_to_stablehlo::Bool, transpose_propagate::Symbol, donated_args::Symbol, optimize_then_pad::Bool, cudnn_hlo_optimize::Bool, compile_options::Missing, sync::Bool, no_nan::Bool, raise::Bool, shardy_passes::Symbol, optimize::Bool, optimize_communications::Bool})
@ Reactant.Compiler ~/.julia/packages/Reactant/gBXlB/src/Compiler.jl:3505
[38] top-level scope
@ ~/.julia/packages/Reactant/gBXlB/src/Compiler.jl:2586
in expression starting at /mnt/c/Users/yolha/Desktop/juju_tests/mini/test/main2.jl:10