allow selecting dimensions with predicates (#618)

Author: rafaqz

src/Dimensions/primitives.jl

@@ -206,10 +206,10 @@ julia> dimnum(A, Y)
     all(hasdim(x, q1, query...)) || _extradimserror()
     _call_primitive(_dimnum, MaybeFirst(), x, q1, query...)
 end
+@inline dimnum(x, query::Function) =
+    _call_primitive(_dimnum, MaybeFirst(), x, query)
 
-@inline function _dimnum(f::Function, ds::Tuple, query::Tuple{Vararg{Int}})
-    query
-end
+@inline _dimnum(f::Function, ds::Tuple, query::Tuple{Vararg{Int}}) = query
 @inline function _dimnum(f::Function, ds::Tuple, query::Tuple)
     numbered = map(ds, ntuple(identity, length(ds))) do d, i
         rebuild(d, i)
@@ -659,6 +659,12 @@ struct AlwaysTuple end
 @inline _call_primitive1(f, t, op::Function, x, query) = _call_primitive1(f, t, op, dims(x), query)
 @inline _call_primitive1(f, t, op::Function, x::Nothing) = _dimsnotdefinederror()
 @inline _call_primitive1(f, t, op::Function, x::Nothing, query) = _dimsnotdefinederror()
+@inline function _call_primitive1(f, t, op::Function, ds::Tuple, query::Function) 
+    selection = foldl(ds; init=()) do acc, d
+        query(d) ? (acc..., d) : acc
+    end
+    _call_primitive1(f, t, op, ds, selection)
+end
 @inline function _call_primitive1(f, t, op::Function, d::Tuple, query)
     ds = dims(query)
     isnothing(ds) && _dims_are_not_dims()

test/primitives.jl

@@ -139,6 +139,8 @@ end
     @test (@ballocated $f1($dimz)) == 0
 
     @test dims(da, X()) isa X
+    @test dims(da, isforward) isa Tuple{<:X,<:Y}
+    @test dims(da, !isforward) isa Tuple{}
     @test dims(da, Z()) isa Nothing
     @test (@inferred dims(da, XDim, YDim)) isa Tuple{<:X,<:Y}
     @test (@ballocated dims($da, XDim, YDim)) == 0
@@ -171,6 +173,7 @@ end
 
 @testset "commondims" begin
     @test commondims(da, X) == (dims(da, X),)
+    @test commondims(da, x -> x isa X) == (dims(da, X),)
     # Dims are always in the base order
     @test (@inferred commondims(da, (Y(), X()))) == dims(da, (X, Y))
     @test (@ballocated commondims($da, (Y(), X()))) == 0
@@ -205,6 +208,7 @@ end
 @testset "dimnum" begin
     dims(da)
     @test dimnum(da, Y()) == dimnum(da, 2) == 2
+    @test dimnum(da, Base.Fix2(isa,Y)) == (2,)
     @test (@ballocated dimnum($da, Y())) == 0
     @test dimnum(da, X) == 1
     @test (@ballocated dimnum($da, X)) == 0
@@ -228,6 +232,7 @@ end
 
 @testset "hasdim" begin
     @test hasdim(da, X()) == true
+    @test hasdim(da, isforward) == (true, true) 
     @test (@ballocated hasdim($da, X())) == 0
     @test hasdim(da, Ti) == false
     @test (@ballocated hasdim($da, Ti)) == 0
@@ -264,6 +269,7 @@ end
 @testset "otherdims" begin
     A = DimArray(ones(5, 10, 15), (X, Y, Z));
     @test otherdims(A, X()) == dims(A, (Y, Z))
+    @test otherdims(A, x -> x isa X) == dims(A, (Y, Z))
     @test (@ballocated otherdims($A, X())) == 0
     @test otherdims(A, Y) == dims(A, (X, Z))
     @test otherdims(A, Z) == dims(A, (X, Y))