Redirect broadcasted setindex via fill!

[avik-pal]

currently, it is legal to do:

fn(x) = begin
     y = copy(x)
     UpperTriangular(y) .= 1.0
     return parent(y)
end

in Reactant. This technically needs #369 to work but that PR is not adding the setindex functionality.

[avik-pal]

We could technically throw an error in set_mlir_data!

[wsmoses]

alternatively within set_mlir_data we could do the corresponding fixes. For example here, we would say (in psedocode)

function set_mlir_data!(x::UpperTriangular{Traced}, data)
    data = Ops.concat_upper_triangular_with_lower_of(data, get_mlir_data(x))
    set_mlir_data!(x.parent, data)
end

[avik-pal]

The above is not legal in Julia since UpperTriangular can't have non-zero entries in the lower triangular part. But the way we implement set_mlir_data! it allows this in Reactant currently.

[wsmoses]

julia> x = ones(4,4)
4×4 Matrix{Float64}:
 1.0  1.0  1.0  1.0
 1.0  1.0  1.0  1.0
 1.0  1.0  1.0  1.0
 1.0  1.0  1.0  1.0

julia> y = UpperTriangular(x)
ERROR: UndefVarError: `UpperTriangular` not defined
Stacktrace:
 [1] top-level scope
   @ REPL[2]:1

julia> using LinearAlgebra

julia> y = UpperTriangular(x)
4×4 UpperTriangular{Float64, Matrix{Float64}}:
 1.0  1.0  1.0  1.0
  ⋅   1.0  1.0  1.0
  ⋅    ⋅   1.0  1.0
  ⋅    ⋅    ⋅   1.0


julia> y .*= 2
4×4 UpperTriangular{Float64, Matrix{Float64}}:
 4.0  2.0  2.0  2.0
  ⋅   2.0  2.0  2.0
  ⋅    ⋅   2.0  2.0
  ⋅    ⋅    ⋅   2.0

julia> x
4×4 Matrix{Float64}:
 4.0  2.0  2.0  2.0
 1.0  2.0  2.0  2.0
 1.0  1.0  2.0  2.0
 1.0  1.0  1.0  2.0

[wsmoses]

I agree structurally it doesn't have non-zero entries there, but the parent itself actually might, so in order to store the data into the parent (the only legal possibility), then we have to fuse the existing entries, with the new ones

[avik-pal]

All operations expect direct setindexing is allowed:

julia> y .= 1
ERROR: ArgumentError: cannot set indices in the lower triangular part of an UpperTriangular matrix to a nonzero value (1)
Stacktrace:
 [1] throw_setindex_structuralzero_error(T::Type, x::Any)
   @ LinearAlgebra ~/.julia/juliaup/julia-1.11.2+0.x64.linux.gnu/share/julia/stdlib/v1.11/LinearAlgebra/src/triangular.jl:347
 [2] fill!
   @ ~/.julia/juliaup/julia-1.11.2+0.x64.linux.gnu/share/julia/stdlib/v1.11/LinearAlgebra/src/triangular.jl:352 [inlined]
 [3] copyto!
   @ ./broadcast.jl:928 [inlined]
 [4] materialize!
   @ ./broadcast.jl:878 [inlined]
 [5] materialize!(dest::UpperTriangular{Float64, Matrix{Float64}}, bc::Base.Broadcast.Broadcasted{Base.Broadcast.DefaultArrayStyle{0}, Nothing, typeof(identity), Tuple{Int64}})
   @ Base.Broadcast ./broadcast.jl:875
 [6] top-level scope
   @ REPL[7]:1
 [7] top-level scope
   @ none:1

[avik-pal]

That is why, I feel redirecting the setindex via fill! is a nicer solution. That way other broadcasted operations continue working

[wsmoses]

Fair enough, (and I have no strong opinions here). Go for it!