Fix some tests

Author: mtfishman

Project.toml

@@ -12,6 +12,7 @@ GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MapBroadcast = "ebd9b9da-f48d-417c-9660-449667d60261"
 MatrixAlgebraKit = "6c742aac-3347-4629-af66-fc926824e5e4"
+TypeParameterAccessors = "7e5a90cf-f82e-492e-a09b-e3e26432c138"
 
 [weakdeps]
 BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e"
@@ -37,4 +38,5 @@ MapBroadcast = "0.1.10"
 MatrixAlgebraKit = "0.2, 0.3"
 TensorAlgebra = "0.3.10"
 TensorProducts = "0.1.7"
+TypeParameterAccessors = "0.4.2"
 julia = "1.10"

ext/KroneckerArraysBlockSparseArraysExt/KroneckerArraysBlockSparseArraysExt.jl

@@ -36,8 +36,7 @@ end
 
 using BlockArrays: AbstractBlockedUnitRange
 using BlockSparseArrays: Block, ZeroBlocks, eachblockaxis, mortar_axis
-using KroneckerArrays: KroneckerArrays, KroneckerArray, ⊗, arg1, arg2
-using BlockSparseArrays.TypeParameterAccessors: unwrap_array_type
+using KroneckerArrays: KroneckerArrays, KroneckerArray, ⊗, arg1, arg2, isactive
 
 function KroneckerArrays.arg1(r::AbstractBlockedUnitRange)
   return mortar_axis(arg1.(eachblockaxis(r)))
@@ -59,28 +58,21 @@ end
 
 ## TODO: Is this needed?
 function Base.getindex(
-  a::ZeroBlocks{N,KroneckerArray{T,N,A,B}}, I::Vararg{Int,N}
-) where {T,N,A<:AbstractArray{T,N},B<:AbstractArray{T,N}}
+  a::ZeroBlocks{N,KroneckerArray{T,N,A1,A2}}, I::Vararg{Int,N}
+) where {T,N,A1<:AbstractArray{T,N},A2<:AbstractArray{T,N}}
   ax_a1 = map(arg1, a.parentaxes)
   ax_a2 = map(arg2, a.parentaxes)
-  # TODO: Instead of mutability, maybe have a trait like
-  # `isstructural` or `isdata`.
-  ismut1 = ismutabletype(unwrap_array_type(A))
-  ismut2 = ismutabletype(unwrap_array_type(B))
-  (ismut1 || ismut2) || error("Can't get zero block.")
-  a1 = if ismut1
-    ZeroBlocks{N,A}(ax_a1)[I...]
-  else
-    block_ax_a1 = arg1.(block_axes(a.parentaxes, Block(I)))
-    similar(A, block_ax_a1)
+  block_ax_a1 = arg1.(block_axes(a.parentaxes, Block(I)))
+  block_ax_a2 = arg2.(block_axes(a.parentaxes, Block(I)))
+  # TODO: Is this a good definition? It is similar to
+  # the definition of `similar` and `adapt_structure`.
+  return if isactive(A1) == isactive(A2)
+    ZeroBlocks{N,A1}(ax_a1)[I...] ⊗ ZeroBlocks{N,A2}(ax_a2)[I...]
+  elseif isactive(A1)
+    ZeroBlocks{N,A1}(ax_a1)[I...] ⊗ A2(block_ax_a2)
+  elseif isactive(A2)
+    A1(block_ax_a1) ⊗ ZeroBlocks{N,A2}(ax_a2)[I...]
   end
-  a2 = if ismut2
-    ZeroBlocks{N,B}(ax_a2)[I...]
-  else
-    block_ax_a2 = arg2.(block_axes(a.parentaxes, Block(I)))
-    a2 = similar(B, block_ax_a2)
-  end
-  return a1 ⊗ a2
 end
 
 using BlockSparseArrays: BlockSparseArrays

src/kroneckerarray.jl

@@ -1,10 +1,30 @@
+# TODO: Move this to DiagonalArrays.jl.
+using DiagonalArrays: DiagonalArrays, _DiagonalArray, DiagonalArray, Unstored
+# TODO: Also support size inputs.
+function DiagonalArrays.DiagonalArray{T,N,D,U}(
+  ax::Tuple{AbstractUnitRange{<:Integer},Vararg{AbstractUnitRange{<:Integer}}}
+) where {T,N,D<:AbstractVector{T},U<:AbstractArray{T,N}}
+  # TODO: Support these constructors.
+  # return DiagonalArray{T,N,Diag,Unstored}(Diag((Base.OneTo(minimum(length, ax)),)), Unstored(U(ax)))
+  # return DiagonalArray{T,N,Diag}(Diag((Base.OneTo(minimum(length, ax)),)), Unstored(U(ax)))
+  # return DiagonalArray{T,N}(Diag((Base.OneTo(minimum(length, ax)),)), Unstored(U(ax)))
+  # return DiagonalArray{T}(D((Base.OneTo(minimum(length, ax)),)), Unstored(U(ax)))
+  # return DiagonalArray(D((Base.OneTo(minimum(length, ax)),)), Unstored(U(ax)))
+  return _DiagonalArray(D((Base.OneTo(minimum(length, ax)),)), U(ax))
+end
+
 function unwrap_array(a::AbstractArray)
   p = parent(a)
   p ≡ a && return a
   return unwrap_array(p)
 end
 isactive(a::AbstractArray) = ismutable(unwrap_array(a))
 
+using TypeParameterAccessors: unwrap_array_type
+function isactive(arrayt::Type{<:AbstractArray})
+  return ismutabletype(unwrap_array_type(arrayt))
+end
+
 # Custom `_convert` works around the issue that
 # `convert(::Type{<:Diagonal}, ::AbstractMatrix)` isn't defined
 # in Julia v1.10 (https://github.com/JuliaLang/julia/pull/48895,
@@ -56,7 +76,17 @@ function mutate_active_args!(f!, f, dest, src)
 end
 
 using Adapt: Adapt, adapt
-Adapt.adapt_structure(to, a::KroneckerArray) = adapt(to, arg1(a)) ⊗ adapt(to, arg2(a))
+function Adapt.adapt_structure(to, a::KroneckerArray)
+  # TODO: Is this a good definition? It is similar to
+  # the definition of `similar`.
+  return if isactive(arg1(a)) == isactive(arg2(a))
+    adapt(to, arg1(a)) ⊗ adapt(to, arg2(a))
+  elseif isactive(arg1(a))
+    adapt(to, arg1(a)) ⊗ arg2(a)
+  elseif isactive(arg2(a))
+    arg1(a) ⊗ adapt(to, arg2(a))
+  end
+end
 
 function Base.copy(a::KroneckerArray)
   return copy(arg1(a)) ⊗ copy(arg2(a))

test/test_blocksparsearrays.jl

@@ -161,10 +161,10 @@ end
   elt in elts
 
   dev = adapt(arrayt)
-  r = @constinferred blockrange([2 × 2, 3 × 3])
+  r = @constinferred blockrange([2 × 2, 2 × 3])
   d = Dict(
     Block(1, 1) => δ(elt, (2, 2)) ⊗ dev(randn(elt, 2, 2)),
-    Block(2, 2) => δ(elt, (3, 3)) ⊗ dev(randn(elt, 3, 3)),
+    Block(2, 2) => δ(elt, (2, 2)) ⊗ dev(randn(elt, 3, 3)),
   )
   a = @constinferred dev(blocksparse(d, (r, r)))
   @test sprint(show, a) == sprint(show, Array(a))
@@ -176,10 +176,10 @@ end
   @test @constinferred(a[Block(2, 2)]) == dev(d[Block(2, 2)])
   @test @constinferred(a[Block(2, 2)]) isa valtype(d)
   @test @constinferred(iszero(a[Block(2, 1)]))
-  @test a[Block(2, 1)] == dev(δ(3, 2) ⊗ zeros(elt, 3, 2))
+  @test a[Block(2, 1)] == dev(δ(2, 2) ⊗ zeros(elt, 3, 2))
   @test a[Block(2, 1)] isa valtype(d)
   @test @constinferred(iszero(a[Block(1, 2)]))
-  @test a[Block(1, 2)] == dev(δ(2, 3) ⊗ zeros(elt, 2, 3))
+  @test a[Block(1, 2)] == dev(δ(2, 2) ⊗ zeros(elt, 2, 3))
   @test a[Block(1, 2)] isa valtype(d)
 
   # Slicing