add missing mul! implementation

.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,12 @@ 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)
+  size(A,2) == size(B,1) || throw(DimensionMismatch("Matrix column must correspond with the OneHot Size $(size(A,2)) ≠ $(size(B,1))")
+)
+  # matmul sometimes wraps in ReshapedArray, taking parent is a simple way to handle that case
+  copyto!(Y, view(A, :, onecold(parent(B))))
+end
+

test/gpu.jl

@@ -26,8 +26,16 @@ end
   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?
+    @test gradient(A -> sum(A * y), gA)[1] isa CuArray
   end
+
+  # 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
+
+  #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

test/runtests.jl

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