Jump.jl: Fixing `axes(::JuMPArray)`

[jrevels]

While tackling #1455, it became apparent that the current indexing behavior of JuMPArray does not conform to the assumptions made by the AbstractArray interface. This will make it quite difficult for JuMPArray to inherit some pretty fundamental features (e.g. broadcast) and compose with other AbstractArray types/operations.

Here are some of Base's indexing assumptions:

CartesianIndices(A) === CartesianIndices(axes(A))
typeof(axes(A, d)) <: AbstractUnitRange{<:Integer}
A[axes(A)...] == A

See https://github.com/JuliaLang/julia/issues/29062 and the Discourse discussion linked therein for more details.

AFAICT, there is no possible way to define axes(::JuMPArray) that a) satisfies all the above assumptions and b) doesn't change user-facing JuMPArray behavior. Right now, axes(::JuMPArray) is overloaded to respect the A[axes(A)...] == A assumption.

My attempt at resolving this without changing user-facing behavior involved defining an AxisRange <: AbstractUnitRange type wrapper, but that ended up being a no-go since Base's R<:AbstractUnitRange code assumes that an instance r::R can be indexed with some multiple of oneunit(eltype(R)) (which doesn't make sense when e.g. eltype(R) === Symbol).

Another way forward is to allow an integer indexing fallback for non-integer JuMPArray axes, e.g.

julia> @variable(m, y[2:5, [:green, :blue]])
2-dimensional JuMPArray{VariableRef,2,...} with index sets:
    Dimension 1, 2:5
    Dimension 2, Symbol[:green, :blue]
And data, a 4×2 Array{VariableRef,2}:
 y[2,green]  y[2,blue]
 y[3,green]  y[3,blue]
 y[4,green]  y[4,blue]
 y[5,green]  y[5,blue]

julia> axes(y)
(2:5, 1:2)

julia> y[3, 2]
y[3,blue]

This approach resolves the underlying issue, but has its own problems. It changes user-facing behavior in some non-trivial ways; previously disallowed indexing would now be allowed, and visa versa. I predict that the main breakage in user code would be when the user specifies an axis to be indexed by a collection with non-concrete eltype, e.g.

julia> @variable(m, x[Any[1, 3, 5], [:green, :blue]])
2-dimensional JuMPArray{VariableRef,2,...} with index sets:
    Dimension 1, Any[1, 3, 5]
    Dimension 2, Symbol[:green, :blue]
And data, a 3×2 Array{VariableRef,2}:
 x[1,green]  x[1,blue]
 x[3,green]  x[3,blue]
 x[5,green]  x[5,blue]

Under this approach, indexing into x above is ambiguous, e.g. it can't be known whether x[3, 2] should return x[3, green] or x[5, green]. It'd probably have to be made an error to construct something like x at all, documenting a user-facing rule like "Given the @variable(m, x[inds...]), all(isconcretetype(eltype(i)) for i in inds) must be true"...gross!

Anybody have any better ideas?

Honestly, it seems like the "right" way to fix this is to just refactor Base's interface, but the ship has sailed on that until Julia 2.0...

[jrevels]

it can't be known whether x[3, 2] should return x[3, green] or x[5, green].

It can't be known from the type information, anyway...

We could just do a runtime lookup for every getindex, and always favor the user-specified index over the fallback, but I was assuming we wouldn't want to do that.

However, we could actually avoid doing the runtime lookup for concrete-eltyped index collections, and only do it for nonconcrete-eltyped index collections. Those getindexs would likely have involved a dynamic dispatch anyway, so maybe it wouldn't be that much slower? Hmm...

[jrevels]

Okay, just thinking out loud here, but another potential solution is to define:

struct FallbackIndex <: Integer
    i::Int
end

and overload getindex(::JuMPArray) on those guys. Then we can define axes(::JuMPArray) to return a tuple of FallbackIndexRange <: AbstractUnitRange{Integer}.

This approach doesn't have the same user-facing problems of the other approach, but it might be harder to implement/more brittle, since it requires defining a new <:Integer subtype...

[mlubin]

AxisArrays has an identical issue: https://github.com/JuliaArrays/AxisArrays.jl/issues/84.

My initial reaction to the proposal to mix different index types like:

julia> y[3, 2]
y[3,blue]

is a strong negative. This is confusing and enables very subtle bugs in user code, and also has the weird edge cases like Any[1, 3, 5] that you brought up. Using a different type to signal integer-based indexing sounds better to me.

[blegat]

I like the FallbackIndex solution, this might also help implement linear indexing through getindex(::JuMPArray, ::FallbackInteger) to fix this test:
https://github.com/JuliaOpt/JuMP.jl/blob/ef17037984d9c10814f80ab5ff90689a4410b5a2/test/variable.jl#L419-L421