Static array shape checking

[aviatesk]

Static array shape checking

[MasonProtter]

Yes, statically sized arrays work out of the box, it's trivial to check:

julia> using StaticArrays, JET

julia> x = SizedVector{3}(rand(3)); y = SizedMatrix{4, 3}(rand(4, 3));

julia> f(x, y) = x' * (y*x)
f (generic function with 1 method)

julia> @report_call f(x, y)
═════ 1 possible error found ═════
┌ @ REPL[9]:1 Main.*(Main.'(x), Main.*(y, x))
│┌ @ /home/mason/julia/usr/share/julia/stdlib/v1.6/LinearAlgebra/src/adjtrans.jl:308 LinearAlgebra.dot(Base.getproperty(u, :parent), v)
││┌ @ /home/mason/.julia/packages/StaticArrays/NTbHj/src/linalg.jl:189 StaticArrays.same_size(a, b)
│││┌ @ /home/mason/.julia/packages/StaticArrays/NTbHj/src/traits.jl:186 StaticArrays._throw_size_mismatch(as...)
││││┌ @ /home/mason/.julia/packages/StaticArrays/NTbHj/src/traits.jl:191 StaticArrays._throw_size_mismatch(::SizedVector{3, Float64, Vector{Float64}}, ::SizedVector{4, Float64, Vector{Float64}})
│││││ may throw: StaticArrays.throw(StaticArrays.DimensionMismatch(Base.string("Sizes ", StaticArrays.map(StaticArrays._size, Argument(2))::Any, " of input arrays do not match")::Any)::Any)
││││└────────────────────────────────────────────────────────────────────
Union{}

julia> @report_call f(x, x*x')
No errors !
Float64

[aviatesk]

Thanks for raising this interesting issue.

In Julia, most information about array shape doesn't appear in Julia's type lattice (i.e. compile time), which its type inference works on, but only appear in runtime.
Since JET.jl is internally powered by Julia's native type inference system (and so it inherits its limitations), JET.jl can reports type level errors like NoMethodError but it can't report those involved with shape and such, like DimensionMismatch.

julia> using JET
[ Info: Precompiling JET [c3a54625-cd67-489e-a8e7-0a5a0ff4e31b]

julia> @report_call [1,2,3]' * [1,2,3]'
═════ 1 possible error found ═════
┌ @ /Users/aviatesk/julia/julia/usr/share/julia/stdlib/v1.7/LinearAlgebra/src/adjtrans.jl:317 *(::LinearAlgebra.Adjoint{Int64, Vector{Int64}}, ::LinearAlgebra.Adjoint{Int64, Vector{Int64}})
│ may throw: LinearAlgebra.throw(LinearAlgebra.MethodError(LinearAlgebra.*, Core.tuple(Argument(2), Argument(3))::Tuple{LinearAlgebra.Adjoint{Int64, Vector{Int64}}, LinearAlgebra.Adjoint{Int64, Vector{Int64}}})::Any)
└─────────────────────────────────────────────────────────────────────────────────────────────
Union{}

julia> @report_call [1,2,3]' * [1,2,3,4]
No errors !
Int64

Static array shape checking will require lots of work, extending the type lattice and maybe huge modification of its abstract interpretation logic, as such I don't think I will working on it in near future (I'm focusing on improving the quality of type-level checks for now).

But, hopefully, once I've got satisfied with the type-level error check, I think a more advanced error check like array shape checking might be my next scope.

---

On the other hand, there are static analysis approaches like : https://github.com/google-research/swift-tfp

Well, JET.jl is actually supposed to check errors mostly statically, and so swift-tfp's approach (abstract interpretation) is in fact same as JET for array shape checking.

Anyway, thanks for telling me interesting sets of discussions and projects !

[AriMKatz]

Glad to hear! What about packages like https://github.com/JuliaArrays/StaticArrays.jl which have array dims in type parameters? (If we can separate codegen from type logic and dispatch)

There are also abstract tracing type approaches : FluxML/Mjolnir.jl#1

Also can you elaborate on the "mostly statically" and the connection to tfp?

Thanks !

[AriMKatz]

That's awesome. Now we just need to deal with the codegen issue right ?

[MasonProtter]

What codegen issue?

[AriMKatz]

JuliaLang/julia#11339

Would make static arrays more versatile