Fix some type ambiguities (#585)

Author: christiangnrd

lib/mps/random.jl

@@ -83,6 +83,7 @@ function Random.rand!(rng::RNG, A::AbstractArray{T, N}) where {T <: Union{Unifor
     end
     return A
 end
+
 function Random.randn!(rng::RNG, A::AbstractArray{T, N}) where {T <: Float32, N}
     isempty(A) && return A
     if MTL.can_alloc_nocopy(pointer(A), sizeof(A))

lib/mpsgraphs/execution.jl

@@ -1,6 +1,6 @@
 
-MPS.encode!(commandBuffer::MPSCommandBuffer, graph::MPSGraph, feeds::MPSGraphTensorDataDictionary, resultsDictionary::MPSGraphTensorDataDictionary) = @inline MPS.encode!(commandBuffer, graph, feeds, nil, resultsDictionary, MPSGraphExecutionDescriptor())
-function MPS.encode!(commandBuffer::MPSCommandBuffer, graph::MPSGraph, feeds::MPSGraphTensorDataDictionary, targetOperations, resultsDictionary::MPSGraphTensorDataDictionary, executionDescriptor)
+MPS.encode!(commandBuffer::MPSCommandBuffer, graph::MPSGraph, feeds::MPSGraphTensorDataDictionary, resultsDictionary::MPSGraphTensorDataDictionary) = @inline MPS.encode!(commandBuffer, graph, feeds, resultsDictionary, nil, MPSGraphExecutionDescriptor())
+function MPS.encode!(commandBuffer::MPSCommandBuffer, graph::MPSGraph, feeds::MPSGraphTensorDataDictionary, resultsDictionary::MPSGraphTensorDataDictionary, targetOperations, executionDescriptor::MPSGraphExecutionDescriptor)
     @objc [graph::id{MPSGraph} encodeToCommandBuffer:commandBuffer::id{MPSCommandBuffer}
                                                   feeds:feeds::id{MPSGraphTensorDataDictionary}
                                        targetOperations:targetOperations::id{Object}
@@ -9,7 +9,7 @@ function MPS.encode!(commandBuffer::MPSCommandBuffer, graph::MPSGraph, feeds::MP
     return resultsDictionary
 end
 
-function MPS.encode!(commandBuffer::MPSCommandBuffer, graph::MPSGraph, feeds::MPSGraphTensorDataDictionary, targetTensors::NSArray, targetOperations=nil, executionDescriptor=MPSGraphExecutionDescriptor())
+function MPS.encode!(commandBuffer::MPSCommandBuffer, graph::MPSGraph, feeds::MPSGraphTensorDataDictionary, targetTensors::NSArray, targetOperations=nil, executionDescriptor::MPSGraphExecutionDescriptor=MPSGraphExecutionDescriptor())
     obj = @objc [graph::id{MPSGraph} encodeToCommandBuffer:commandBuffer::id{MPSCommandBuffer}
                                                         feeds:feeds::id{MPSGraphTensorDataDictionary}
                                                 targetTensors:targetTensors::id{NSArray}

lib/mpsgraphs/matmul.jl

@@ -86,7 +86,7 @@ end
     )
 
     cmdbuf = MPSCommandBuffer(Metal.global_queue(device()))
-    encode!(cmdbuf, graph, NSDictionary(feeds), nil, NSDictionary(resultdict), default_exec_desc())
+    encode!(cmdbuf, graph, NSDictionary(feeds), NSDictionary(resultdict), nil, default_exec_desc())
     commit!(cmdbuf)
     wait_completed(cmdbuf)
 

src/indexing.jl

@@ -52,9 +52,12 @@ function Base.findall(bools::WrappedMtlArray{Bool})
     return ys
 end
 
-function Base.findall(f::Function, A::WrappedMtlArray)
+@inline function _findall(f, A)
     bools = map(f, A)
     ys = findall(bools)
     unsafe_free!(bools)
     return ys
-end
+end
+
+Base.findall(f::Function, A::WrappedMtlArray) = _findall(f, A)
+Base.findall(f::Base.Fix2{typeof(in)}, A::WrappedMtlArray) = _findall(f, A)

test/array.jl

@@ -540,6 +540,12 @@ end
     let x = rand(Float32, 1000, 1000)
         @test findall(y->y>Float32(0.5), x) == Array(findall(y->y>Float32(0.5), MtlArray(x)))
     end
+
+    # ambiguity
+    let f = in(3)
+        x = MtlArray([1, 2, 3, 4, 5, 3])
+        @test Array(findall(f, x)) == [3, 6]
+    end
 end
 
 @testset "broadcast" begin