add missing mul! implementation (#47)

* add missing mul! implementation

* remove broken from working test

* bump version

* add broken test

* remove test_broken note that is no longer relevant

* LTS is now 1.10

* bump buildkite julia version

* clean up tests

* check dimensions of output matrix in mul!

* Update linalg.jl

* Update linalg.jl

* Update src/linalg.jl

Co-authored-by: Michael Abbott <[email protected]>

---------

Co-authored-by: Michael Abbott <[email protected]>

Author: ExpandingMan

.buildkite/pipeline.yml

@@ -1,11 +1,11 @@
 steps:
-  - label: "GPU integration with julia v1.6"
+  - label: "GPU integration with julia v1.10"
     plugins:
       - JuliaCI/julia#v1:
           # Drop default "registries" directory, so it is not persisted from execution to execution
           # Taken from https://github.com/JuliaLang/julia/blob/v1.7.2/.buildkite/pipelines/main/platforms/package_linux.yml#L11-L12
           persist_depot_dirs: packages,artifacts,compiled
-          version: "1.6"
+          version: "1.10"
       - JuliaCI/julia-test#v1: ~
     agents:
       queue: "juliagpu"

.github/workflows/CI.yml

@@ -18,7 +18,7 @@ jobs:
       fail-fast: false
       matrix:
         version:
-          - '1.6'
+          - '1.10'
           - '1'
           - 'nightly'
         os:
@@ -47,17 +47,17 @@ jobs:
 
       - name: "Run test without coverage report"
         uses: julia-actions/julia-runtest@v1
-        if: ${{ !contains(fromJson('["1", "1.6"]'), matrix.version) || matrix.os != 'ubuntu-latest' }}
+        if: ${{ !contains(fromJson('["1", "1.10"]'), matrix.version) || matrix.os != 'ubuntu-latest' }}
         with:
           coverage: false
 
       - name: "Run test with coverage report"
         uses: julia-actions/julia-runtest@v1
-        if: contains(fromJson('["1", "1.6"]'), matrix.version) && matrix.os == 'ubuntu-latest'
+        if: contains(fromJson('["1", "1.10"]'), matrix.version) && matrix.os == 'ubuntu-latest'
       - uses: julia-actions/julia-processcoverage@v1
-        if: contains(fromJson('["1", "1.6"]'), matrix.version) && matrix.os == 'ubuntu-latest'
+        if: contains(fromJson('["1", "1.10"]'), matrix.version) && matrix.os == 'ubuntu-latest'
       - uses: codecov/codecov-action@v3
-        if: contains(fromJson('["1", "1.6"]'), matrix.version) && matrix.os == 'ubuntu-latest'
+        if: contains(fromJson('["1", "1.10"]'), matrix.version) && matrix.os == 'ubuntu-latest'
         with:
           files: lcov.info
 
@@ -68,7 +68,7 @@ jobs:
       - uses: actions/checkout@v3
       - uses: julia-actions/setup-julia@v1
         with:
-          version: '1.6'
+          version: '1.10'
       - run: |
           julia --project=docs -e '
             using Pkg

Project.toml

@@ -1,6 +1,6 @@
 name = "OneHotArrays"
 uuid = "0b1bfda6-eb8a-41d2-88d8-f5af5cad476f"
-version = "0.2.6"
+version = "0.2.7"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

src/linalg.jl

@@ -33,3 +33,20 @@ for wrapper in [:Adjoint, :Transpose]
     end
   end
 end
+
+function LinearAlgebra.mul!(Y::AbstractVecOrMat, A::AbstractMatrix, B::OneHotLike)
+  _isonehot(B) || return invoke(mul!, Tuple{AbstractArray,AbstractMatrix,AbstractMatrix}, Y, A, B)
+  if size(A,2) ≠ size(B,1)
+    throw(DimensionMismatch("Matrix column must correspond with the OneHot Size $(size(A,2)) ≠ $(size(B,1))"))
+  end
+  if !(size(Y,1) == size(A,1) && size(Y,2) == size(B,2))
+    throw(DimensionMismatch("Invalid output matrix size for multiplication of matrix sizes $(size(A)) and $(size(B))"))
+  end
+  idxs =  _indices(B)
+  if idxs isa Integer  # occurs whe B is AbstractVector
+    copyto!(Y, view(A, :, idxs))
+  else
+    NNlib.gather!(Y, A, idxs)
+  end
+end
+

test/gpu.jl

@@ -23,11 +23,23 @@ end
   @test (repr("text/plain", y); true)
 
   gA = rand(3, 2) |> cu;
-  if VERSION >= v"1.9" && CUDA.functional()
-    @test gradient(A -> sum(A * y), gA)[1] isa CuArray 
-  else
-    @test_broken gradient(A -> sum(A * y), gA)[1] isa CuArray  # fails with JLArray, bug in Zygote?
-  end
+
+  #NOTE: this would require something that can copute gradient... we don't have that here?
+  #@test gradient(A -> sum(A * y), gA)[1] isa CuArray 
+
+  # some specialized implementations call only mul! and not *, so we must ensure this works
+  @test LinearAlgebra.mul!(similar(gA, 3, 3), gA, y) ≈ gA*y
+  @test LinearAlgebra.mul!(similar(gA, 3, 1), gA, onehot(1, 1:2)) ≈ gA*onehot(1, 1:2)
+
+  @test_throws DimensionMismatch LinearAlgebra.mul!(similar(gA, 3, 4), gA, y)
+
+  gB = rand(3, 3) |> cu
+  @test_throws DimensionMismatch LinearAlgebra.mul!(similar(gB, 3, 3), gB, y)
+
+  #TODO: the below fails due to method ambiguity and GPU scalar indexing
+  y = reshape(y, 3, 2)
+  gA = rand(2, 3) |> cu
+  @test_broken LinearAlgebra.mul!(similar(gA, 2, 2), gA, y) ≈ gA*y
 end
 
 @testset "onehotbatch(::CuArray, ::UnitRange)" begin
@@ -48,7 +60,9 @@ end
   y = onehotbatch(ones(3), 1:10) |> cu;
   l = ['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j']
   @test onecold(y) isa CuArray
+  @test onecold(y) == cu([1, 1, 1])  # == doesn't work across devices
   @test y[3,:] isa CuArray
+  @test y[3,:] == cu([0, 0, 0])  # == doesn't work across devices
   @test onecold(y, l) == ['a', 'a', 'a']
 end
 

test/linalg.jl

@@ -12,6 +12,17 @@
   @test transpose(X) * b2 == transpose(X) * Array(b2)
   @test transpose(v) * b2 == transpose(v) * Array(b2)
   @test_throws DimensionMismatch A*b2
+
+  # in-place with mul!
+  c1 = fill(NaN, 3)
+  @test mul!(c1, A, b1) == A * Array(b1)
+  @test c1 == A * Array(b1)
+  @test mul!(c1, transpose(A), b1) == transpose(A) * Array(b1)
+  @test mul!(zeros(3,1), A, b1) == reshape(A * b1, 3,1)
+  @test mul!([NaN], transpose(v), b2) == mul!([NaN], transpose(v), Array(b2))
+  
+  @test_throws DimensionMismatch mul!(zeros(5), A, b1)
+  @test_throws DimensionMismatch mul!(c1, X, b1)
 end
 
 @testset "AbstractMatrix-OneHotMatrix multiplication" begin
@@ -22,8 +33,10 @@ end
   b2 = OneHotMatrix([2, 4, 1, 3], 5)
   b3 = OneHotMatrix([1, 1, 2], 4)
   b4 = reshape(OneHotMatrix([1 2 3; 2 2 1], 3), 3, :)
+  @test OneHotArrays._isonehot(b4)
   b5 = reshape(b4, 6, :)
   b6 = reshape(OneHotMatrix([1 2 2; 2 2 1], 2), 3, :)
+  @test !OneHotArrays._isonehot(b6)
   b7 = reshape(OneHotMatrix([1 2 3; 1 2 3], 3), 6, :)
 
   @test A * b1 == A[:,[1, 1, 2, 2]]
@@ -41,4 +54,21 @@ end
   @test_throws DimensionMismatch A*b2'
   @test_throws DimensionMismatch A*b6'
   @test_throws DimensionMismatch A*b7
+
+  # in-place with mul!
+  c1 = fill(NaN, 3, 4)
+  @test mul!(c1, A, b1) == A * b1
+  @test c1 == A * b1
+  
+  c4 = fill(NaN, 3, 6)
+  @test mul!(c4, A, b4) == A * b4  # b4 is reshaped but still one-hot
+  @test mul!(c4, A', b4) == A' * b4
+  c6 = fill(NaN, 3, 4)
+  @test mul!(c6, A, b6) == A * b6  # b4 is reshaped and not one-hot
+  @test mul!(c6, A', b6) == A' * b6
+  
+  @test_throws DimensionMismatch mul!(c1, A, b2)
+  @test_throws DimensionMismatch mul!(c1, A, b4)
+  @test_throws DimensionMismatch mul!(c4, A, b1)
+  @test_throws DimensionMismatch mul!(zeros(10, 3), A, b1)
 end

test/runtests.jl

@@ -1,5 +1,5 @@
 using OneHotArrays
-using Test
+using Test, LinearAlgebra
 using Compat: stack
 
 @testset "OneHotArray" begin