Infer static size of subarrays of adjoint vectors

Author: chriselrod

src/axes.jl

@@ -48,11 +48,16 @@ function axes_types(::Type{T}) where {T<:AbstractRange}
     end
 end
 axes_types(::Type{T}) where {T<:ReshapedArray} = NTuple{ndims(T),OneTo{Int}}
-function _sub_axis_type(::Type{I}, dim::StaticInt{D}) where {I<:Tuple,D}
-    axes_types(field_type(I, dim), static(1))
+function _sub_axis_type(::Type{I}, ::Type{PA}, dim::StaticInt{D}) where {I<:Tuple,PA,D}
+    IT = field_type(I, dim)
+    if IT <: Base.Slice
+        axes_types(field_type(PA, dim), static(1))
+    else
+        axes_types(IT, static(1))
+    end
 end
 @inline function axes_types(::Type{T}) where {N,P,I,T<:SubArray{<:Any,N,P,I}}
-    return eachop_tuple(_sub_axis_type, to_parent_dims(T), I)
+    return eachop_tuple(_sub_axis_type, to_parent_dims(T), axes_types(P), I)
 end
 
 function axes_types(::Type{T}) where {T<:ReinterpretArray}

test/axes.jl

@@ -67,7 +67,14 @@ end
     @test @inferred(ArrayInterface.axes(b, static(3))) == 1:1
 end
 
+@testset "SubArray Adjoint Axis" begin
+  N = 4; d = rand(N);
+
+  @test @inferred(ArrayInterface.axes_types(typeof(view(d',:,1:2)))) === Tuple{ArrayInterface.OptionallyStaticUnitRange{Static.StaticInt{1}, Static.StaticInt{1}}, Base.OneTo{Int64}}
+
+end
 if isdefined(Base, :ReshapedReinterpretArray)
+  @testset "ReshapedReinterpretArray" begin
     a = rand(3, 5)
     ua = reinterpret(reshape, UInt64, a)
     @test ArrayInterface.axes(ua) === ArrayInterface.axes(a)
@@ -79,4 +86,5 @@ if isdefined(Base, :ReshapedReinterpretArray)
     @test @inferred(ArrayInterface.axes(u8a, static(2))) isa ArrayInterface.axes_types(u8a, 2)
     fa = reinterpret(reshape, Float64, copy(u8a))
     @inferred(ArrayInterface.axes(fa)) isa ArrayInterface.axes_types(fa)
+  end
 end