More general scalar indexing

Author: mtfishman

src/kroneckerarray.jl

@@ -158,26 +158,13 @@ function Base.getindex(a::KroneckerArray, i::Integer)
   return a[CartesianIndices(a)[i]]
 end
 
-# TODO: Use this logic from KroneckerProducts.jl for cartesian indexing
-# in the n-dimensional case and use it to replace the matrix and vector cases:
-# https://github.com/perrutquist/KroneckerProducts.jl/blob/8c0104caf1f17729eb067259ba1473986121d032/src/KroneckerProducts.jl#L59-L66
-function Base.getindex(a::KroneckerArray{<:Any,N}, I::Vararg{Integer,N}) where {N}
-  return error("Not implemented.")
-end
-
 using GPUArraysCore: GPUArraysCore
-function Base.getindex(a::KroneckerMatrix, i1::Integer, i2::Integer)
-  GPUArraysCore.assertscalar("getindex")
-  # Code logic from Kronecker.jl:
-  # https://github.com/MichielStock/Kronecker.jl/blob/v0.5.5/src/base.jl#L101-L105
-  k, l = size(arg2(a))
-  return arg1(a)[cld(i1, k), cld(i2, l)] * arg2(a)[(i1 - 1) % k + 1, (i2 - 1) % l + 1]
-end
-
-function Base.getindex(a::KroneckerVector, i::Integer)
+function Base.getindex(a::KroneckerArray{<:Any,N}, I::Vararg{Integer,N}) where {N}
   GPUArraysCore.assertscalar("getindex")
-  k = length(arg2(a))
-  return arg1(a)[cld(i, k)] * arg2(a)[(i - 1) % k + 1]
+  I′ = ntuple(Val(N)) do dim
+    return cartesianproduct(axes(a, dim))[I[dim]]
+  end
+  return a[I′...]
 end
 
 # Allow customizing for `FillArrays.Eye`.

test/Project.toml

@@ -6,6 +6,7 @@ BlockSparseArrays = "2c9a651f-6452-4ace-a6ac-809f4280fbb4"
 DerivableInterfaces = "6c5e35bf-e59e-4898-b73c-732dcc4ba65f"
 DiagonalArrays = "74fd4be6-21e2-4f6f-823a-4360d37c7a77"
 FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
+GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527"
 JLArrays = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb"
 KroneckerArrays = "05d0b138-81bc-4ff7-84be-08becefb1ccc"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
@@ -24,6 +25,7 @@ BlockSparseArrays = "0.7.19"
 DerivableInterfaces = "0.5"
 DiagonalArrays = "0.3.7"
 FillArrays = "1"
+GPUArraysCore = "0.2"
 JLArrays = "0.2"
 KroneckerArrays = "0.1"
 LinearAlgebra = "1.10"

test/test_basics.jl

@@ -2,6 +2,7 @@ using Adapt: adapt
 using Base.Broadcast: BroadcastStyle, Broadcasted, broadcasted
 using DerivableInterfaces: zero!
 using DiagonalArrays: diagonal
+using GPUArraysCore: @allowscalar
 using JLArrays: JLArray
 using KroneckerArrays:
   KroneckerArrays,
@@ -44,6 +45,15 @@ elts = (Float32, Float64, ComplexF32, ComplexF64)
   @test first(r) == 2
   @test last(r) == 7
 
+  # Test high-dimensional materialization.
+  a = randn(elt, 2, 2, 2) ⊗ randn(elt, 2, 2, 2)
+  x = Array(a)
+  y = similar(x)
+  for I in eachindex(a)
+    y[I] = @allowscalar x[I]
+  end
+  @test x == y
+
   a = @constinferred(randn(elt, 2, 2) ⊗ randn(elt, 3, 3))
   b = randn(elt, 2, 2) ⊗ randn(elt, 3, 3)
   c = a.a ⊗ b.b