|
1 | 1 |
|
| 2 | +#julia> @btime ArrayInterface.is_increasing(ArrayInterface.nstatic(Val(10))) |
| 3 | +# 0.045 ns (0 allocations: 0 bytes) |
| 4 | +#ArrayInterface.True() |
| 5 | +function is_increasing(perm::Tuple{StaticInt{X},StaticInt{Y},Vararg}) where {X, Y} |
| 6 | + if X <= Y |
| 7 | + return is_increasing(tail(perm)) |
| 8 | + else |
| 9 | + return False() |
| 10 | + end |
| 11 | +end |
| 12 | +function is_increasing(perm::Tuple{StaticInt{X},StaticInt{Y}}) where {X, Y} |
| 13 | + if X <= Y |
| 14 | + return True() |
| 15 | + else |
| 16 | + return False() |
| 17 | + end |
| 18 | +end |
| 19 | +is_increasing(::Tuple{StaticInt{X}}) where {X} = True() |
| 20 | + |
| 21 | +""" |
| 22 | + from_parent_dims(::Type{T}) -> Bool |
| 23 | +
|
| 24 | +Returns the mapping from parent dimensions to child dimensions. |
| 25 | +""" |
| 26 | +from_parent_dims(::Type{T}) where {T} = nstatic(Val(ndims(T))) |
| 27 | +from_parent_dims(::Type{T}) where {T<:Union{Transpose,Adjoint}} = (StaticInt(2), One()) |
| 28 | +from_parent_dims(::Type{<:SubArray{T,N,A,I}}) where {T,N,A,I} = _from_sub_dims(A, I) |
| 29 | +@generated function _from_sub_dims(::Type{A}, ::Type{I}) where {A,N,I<:Tuple{Vararg{Any,N}}} |
| 30 | + out = Expr(:tuple) |
| 31 | + n = 1 |
| 32 | + for p in I.parameters |
| 33 | + if argdims(A, p) > 0 |
| 34 | + push!(out.args, :(StaticInt($n))) |
| 35 | + n += 1 |
| 36 | + else |
| 37 | + push!(out.args, :(StaticInt(0))) |
| 38 | + end |
| 39 | + end |
| 40 | + out |
| 41 | +end |
| 42 | +function from_parent_dims(::Type{<:PermutedDimsArray{T,N,<:Any,I}}) where {T,N,I} |
| 43 | + return _val_to_static(Val(I)) |
| 44 | +end |
| 45 | + |
| 46 | +""" |
| 47 | + to_parent_dims(::Type{T}) -> Bool |
| 48 | +
|
| 49 | +Returns the mapping from child dimensions to parent dimensions. |
| 50 | +""" |
| 51 | +to_parent_dims(x) = to_parent_dims(typeof(x)) |
| 52 | +to_parent_dims(::Type{T}) where {T} = nstatic(Val(ndims(T))) |
| 53 | +to_parent_dims(::Type{T}) where {T<:Union{Transpose,Adjoint}} = (StaticInt(2), One()) |
| 54 | +to_parent_dims(::Type{<:PermutedDimsArray{T,N,I}}) where {T,N,I} = _val_to_static(Val(I)) |
| 55 | +to_parent_dims(::Type{<:SubArray{T,N,A,I}}) where {T,N,A,I} = _to_sub_dims(A, I) |
| 56 | +@generated function _to_sub_dims(::Type{A}, ::Type{I}) where {A,N,I<:Tuple{Vararg{Any,N}}} |
| 57 | + out = Expr(:tuple) |
| 58 | + n = 1 |
| 59 | + for p in I.parameters |
| 60 | + if argdims(A, p) > 0 |
| 61 | + push!(out.args, :(StaticInt($n))) |
| 62 | + end |
| 63 | + n += 1 |
| 64 | + end |
| 65 | + out |
| 66 | +end |
| 67 | + |
2 | 68 | """ |
3 | 69 | has_dimnames(::Type{T}) -> Bool |
4 | 70 |
|
|
137 | 203 | end |
138 | 204 | return Expr(:tuple, exs...) |
139 | 205 | end |
| 206 | +@generated function _perm_tuple(::Type{T}, ::Val{P}) where {T,P} |
| 207 | + out = Expr(:curly, :Tuple) |
| 208 | + for p in P |
| 209 | + push!(out.args, T.parameters[p]) |
| 210 | + end |
| 211 | + Expr(:block, Expr(:meta, :inline), out) |
| 212 | +end |
| 213 | + |
| 214 | +""" |
| 215 | + axes_types(::Type{T}[, d]) -> Type |
| 216 | +
|
| 217 | +Returns the type of the axes for `T` |
| 218 | +""" |
| 219 | +axes_types(x) = axes_types(typeof(x)) |
| 220 | +axes_types(x, d) = axes_types(typeof(x), d) |
| 221 | +@inline axes_types(::Type{T}, d) where {T} = axes_types(T).parameters[to_dims(T, d)] |
| 222 | +function axes_types(::Type{T}) where {T} |
| 223 | + if parent_type(T) <: T |
| 224 | + return Tuple{Vararg{OptionallyStaticUnitRange{One,Int},ndims(T)}} |
| 225 | + else |
| 226 | + return axes_types(parent_type(T)) |
| 227 | + end |
| 228 | +end |
| 229 | +function axes_types(::Type{T}) where {T<:Adjoint} |
| 230 | + return _perm_tuple(axes_types(parent_type(T)), Val((2, 1))) |
| 231 | +end |
| 232 | +function axes_types(::Type{T}) where {T<:Transpose} |
| 233 | + return _perm_tuple(axes_types(parent_type(T)), Val((2, 1))) |
| 234 | +end |
| 235 | +function axes_types(::Type{T}) where {I1,T<:PermutedDimsArray{<:Any,<:Any,I1}} |
| 236 | + return _perm_tuple(axes_types(parent_type(T)), Val(I1)) |
| 237 | +end |
| 238 | +function axes_types(::Type{T}) where {T<:AbstractRange} |
| 239 | + if known_length(T) === nothing |
| 240 | + return Tuple{OptionallyStaticUnitRange{One,Int}} |
| 241 | + else |
| 242 | + return Tuple{OptionallyStaticUnitRange{One,StaticInt{known_length(T)}}} |
| 243 | + end |
| 244 | +end |
| 245 | + |
| 246 | +@inline function axes_types(::Type{T}) where {P,I,T<:SubArray{<:Any,<:Any,P,I}} |
| 247 | + return _sub_axes_types(Val(ArrayStyle(T)), I, axes_types(P)) |
| 248 | +end |
| 249 | +@generated function _sub_axes_types( |
| 250 | + ::Val{S}, |
| 251 | + ::Type{I}, |
| 252 | + ::Type{PI}, |
| 253 | +) where {S,I<:Tuple,PI<:Tuple} |
| 254 | + out = Expr(:curly, :Tuple) |
| 255 | + d = 1 |
| 256 | + for i in I.parameters |
| 257 | + ad = argdims(S, i) |
| 258 | + if ad > 0 |
| 259 | + push!(out.args, :(sub_axis_type($(PI.parameters[d]), $i))) |
| 260 | + d += ad |
| 261 | + else |
| 262 | + d += 1 |
| 263 | + end |
| 264 | + end |
| 265 | + Expr(:block, Expr(:meta, :inline), out) |
| 266 | +end |
| 267 | + |
| 268 | +@inline function axes_types(::Type{T}) where {T<:Base.ReinterpretArray} |
| 269 | + return _reinterpret_axes_types( |
| 270 | + axes_types(parent_type(T)), |
| 271 | + eltype(T), |
| 272 | + eltype(parent_type(T)), |
| 273 | + ) |
| 274 | +end |
| 275 | +@generated function _reinterpret_axes_types( |
| 276 | + ::Type{I}, |
| 277 | + ::Type{T}, |
| 278 | + ::Type{S}, |
| 279 | +) where {I<:Tuple,T,S} |
| 280 | + out = Expr(:curly, :Tuple) |
| 281 | + for i = 1:length(I.parameters) |
| 282 | + if i === 1 |
| 283 | + push!(out.args, reinterpret_axis_type(I.parameters[1], T, S)) |
| 284 | + else |
| 285 | + push!(out.args, I.parameters[i]) |
| 286 | + end |
| 287 | + end |
| 288 | + Expr(:block, Expr(:meta, :inline), out) |
| 289 | +end |
| 290 | + |
| 291 | +function axes_types(::Type{T}) where {N,T<:Base.ReshapedArray{<:Any,N}} |
| 292 | + return Tuple{Vararg{OptionallyStaticUnitRange{One,Int},N}} |
| 293 | +end |
| 294 | + |
140 | 295 |
|
141 | 296 | """ |
142 | 297 | size(A) |
|
162 | 317 | return (One(), static_length(x)) |
163 | 318 | end |
164 | 319 |
|
| 320 | +function size(B::S) where {N,NP,T,A<:AbstractArray{T,NP},I,S<:SubArray{T,N,A,I}} |
| 321 | + return _size(size(parent(B)), B.indices, map(static_length, B.indices)) |
| 322 | +end |
| 323 | +function strides(B::S) where {N,NP,T,A<:AbstractArray{T,NP},I,S<:SubArray{T,N,A,I}} |
| 324 | + return _strides(strides(parent(B)), B.indices) |
| 325 | +end |
| 326 | +@generated function _size(A::Tuple{Vararg{Any,N}}, inds::I, l::L) where {N,I<:Tuple,L} |
| 327 | + t = Expr(:tuple) |
| 328 | + for n = 1:N |
| 329 | + if (I.parameters[n] <: Base.Slice) |
| 330 | + push!(t.args, :(@inbounds(_try_static(A[$n], l[$n])))) |
| 331 | + elseif I.parameters[n] <: Number |
| 332 | + nothing |
| 333 | + else |
| 334 | + push!(t.args, Expr(:ref, :l, n)) |
| 335 | + end |
| 336 | + end |
| 337 | + Expr(:block, Expr(:meta, :inline), t) |
| 338 | +end |
| 339 | +@inline size(v::AbstractVector) = (static_length(v),) |
| 340 | +@inline size(B::MatAdjTrans) = permute(size(parent(B)), Val{(2, 1)}()) |
| 341 | +@inline function size(B::PermutedDimsArray{T,N,I1,I2,A}) where {T,N,I1,I2,A} |
| 342 | + return permute(size(parent(B)), Val{I1}()) |
| 343 | +end |
| 344 | +@inline size(A::AbstractArray, ::StaticInt{N}) where {N} = size(A)[N] |
| 345 | +@inline size(A::AbstractArray, ::Val{N}) where {N} = size(A)[N] |
165 | 346 | """ |
166 | 347 | axes(A, d) |
167 | 348 |
|
|
0 commit comments