Julia: Broadcasting UniformScaling Operations

[ChrisRackauckas]

Broadcasting is not only super convenient, but it also gives new features. For example, broadcasted operations can generate GPU kernels with GPUArrays.jl. So it's really nice to be able to broadcast operations... but this one seems missing. I can't seem to do:

      for j in 1:length(u), i in 1:length(u)
          W[i,j] = I[i,j]-J[i,j]
      end

If you try to do something similar with broadcast, it's incorrect:

I .- ones(4,4) # == zeros(4,4)

It would be nice if broadcast worked on the UniformScaling operator.

[tkoolen]

Between this and https://github.com/JuliaLang/julia/issues/23156, is there a need for an identity matrix struct that is a subtype of AbstractMatrix? Something like

struct IdentityMatrix{T} <: AbstractMatrix{T}
    size::Int
end
IdentityMatrix(size::Int) = IdentityMatrix{Float64}(size)
Base.size(A::IdentityMatrix) = (A.size, A.size)
function Base.getindex(A::IdentityMatrix{T}, i::Int, j::Int) where {T}
    @boundscheck checkbounds(A, i, j)
    ifelse(i == j, one(T), zero(T))
end

julia> IdentityMatrix(4) .- zeros(4, 4)
4×4 Array{Float64,2}:
 1.0  0.0  0.0  0.0
 0.0  1.0  0.0  0.0
 0.0  0.0  1.0  0.0
 0.0  0.0  0.0  1.0

[andreasnoack]

I don't think we need to give up on making UniformScaling work as it is. Hopefully, explicit sizes will not be needed. I could hack together a working version fairly easily by defining _containertype(::UniformScaling) and newindexer(shape, ::UniformScaling) methods.

[pabloferz]

What would I .+ 1 be?

[andreasnoack]

Hopefully 2. Would that be difficult to handle?

[pabloferz]

It would make the implementation a bit more cumbersome, but I can give it a shot.

[mbauman]

The current meaning here is now deprecated, so we have space to implement broadcasting as you expected in the future:

julia> I .- ones(4,4)
┌ Warning: broadcast will default to iterating over its arguments in the future. Wrap arguments of
│ type `x::UniformScaling{Bool}` with `Ref(x)` to ensure they broadcast as "scalar" elements.
│   caller = ip:0x0
└ @ Core :-1
4×4 Array{Float64,2}:
 0.0  0.0  0.0  0.0
 0.0  0.0  0.0  0.0
 0.0  0.0  0.0  0.0
 0.0  0.0  0.0  0.0

[mbauman]

Let's re-open this to track the possibility of actually supporting I in broadcast as an array-like object.

[chriscoey]

Bump! Anyone have an idea on how to do this?

[StefanKarpinski]

💯. This came up on Slack today. I wanted to tell someone to write A .+= I to increase the diagonal of A in place, but it doesn't work because A + I fails.

[chriscoey]

bump again. @mbauman any ideas on how to do this?

[mbauman]

I started poking at this earlier this week in https://github.com/JuliaLang/julia/tree/mb/I%2Cbroadcast. That's not quite right; I'm still brainstorming other ways to go about it.

[haampie]

Another bump, apparently this is an issue that often pops up. What is currently blocking?

[StefanKarpinski]

Someone implementing it I think.

[haampie]

I started poking at this earlier this week in https://github.com/JuliaLang/julia/tree/mb/I%2Cbroadcast. That's not quite right; I'm still brainstorming other ways to go about it.

@mbauman do you remember what was not quite right about this?

[mbauman]

Oooh, I don't remember. Should have left better notes for myself. I see I have a TODO about ensuring that the size is square, and looking at the diff now makes me think there may be some holes around heterogeneous combinations with scalar and 0-dim arguments, but I'm not sure.

[francescoalemanno]

bump! :) this comes up often, and kinda defeats the purpose of having a lazy identity matrix

[lindnemi]

+1, i ran into this recently as well.

[vtjnash]

I contemplated taking a slightly different approach, but I can't say if it was necessarily better:

julia> using LinearAlgebra
julia> Base.Broadcast.broadcastable(x::UniformScaling) = x
julia> Base.Broadcast.BroadcastStyle(::Type{<:UniformScaling}) = Base.Broadcast.Unknown()
julia> Base.Broadcast.axes(::UniformScaling) = ()
julia> Base.@propagate_inbounds Base.Broadcast._broadcast_getindex(A::UniformScaling, I) = A[Tuple(I)...]
julia> I .- ones(4,4) # == zeros(4,4)
4×4 Matrix{Float64}:
  0.0  -1.0  -1.0  -1.0
 -1.0   0.0  -1.0  -1.0
 -1.0  -1.0   0.0  -1.0
 -1.0  -1.0  -1.0   0.0

julia> I .+ I
ERROR: ArgumentError: broadcasting requires an assigned BroadcastStyle

julia> 2 .* I
UniformScaling{Int64}
2*I

julia> Bidiagonal([1, 2, 3], [4, 5], :U) .+ I
3×3 Matrix{Int64}:
 2  4  0
 0  3  5
 0  0  4

julia> I .+ ones(5) # TODO: should be a better error?
ERROR: MethodError: no method matching getindex(::UniformScaling{Bool}, ::Int64)

julia> I .+ ones(5,2) # TODO: should be an error
5×2 Matrix{Float64}:
 2.0  1.0
 1.0  2.0
 1.0  1.0
 1.0  1.0
 1.0  1.0

julia> I .+ spzeros(5, 5) # XXX: BAD??
5×5 SparseMatrixCSC{Float64, Int64} with 25 stored entries:
 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  1.0  1.0  1.0  1.0
 1.0  1.0  1.0  1.0  1.0