Fix axes of CartesianProductVector

Author: mtfishman

src/cartesianproduct.jl

@@ -62,7 +62,9 @@ unproduct(r::CartesianProductVector) = getfield(r, :values)
 Base.length(a::CartesianProductVector) = length(unproduct(a))
 Base.size(a::CartesianProductVector) = (length(a),)
 function Base.axes(r::CartesianProductVector)
-  return (CartesianProductUnitRange(cartesianproduct(r), only(axes(unproduct(r)))),)
+  prod = cartesianproduct(r)
+  prod_ax = only(axes(arg1(prod))) × only(axes(arg2(prod)))
+  return (CartesianProductUnitRange(prod_ax, only(axes(unproduct(r)))),)
 end
 function Base.copy(a::CartesianProductVector)
   return CartesianProductVector(copy(cartesianproduct(a)), copy(unproduct(a)))

test/test_basics.jl

@@ -9,6 +9,7 @@ using KroneckerArrays:
   KroneckerArray,
   KroneckerStyle,
   CartesianProductUnitRange,
+  CartesianProductVector,
   ⊗,
   ×,
   arg1,
@@ -45,6 +46,14 @@ elts = (Float32, Float64, ComplexF32, ComplexF64)
   @test r[2 × 2] == 5
   @test r[2 × 3] == 6
 
+  # CartesianProductUnitRange axes
+  r = cartesianrange((2:3) × (3:4), 2:5)
+  @test axes(r) ≡ (CartesianProductUnitRange(Base.OneTo(2) × Base.OneTo(2), Base.OneTo(4)),)
+
+  # CartesianProductVector axes
+  r = CartesianProductVector(([2, 4]) × ([3, 5]), [3, 5, 7, 9])
+  @test axes(r) ≡ (CartesianProductUnitRange(Base.OneTo(2) × Base.OneTo(2), Base.OneTo(4)),)
+
   r = @constinferred(cartesianrange(2 × 3, 2:7))
   @test r === cartesianrange(Base.OneTo(2) × Base.OneTo(3), 2:7)
   @test cartesianproduct(r) === Base.OneTo(2) × Base.OneTo(3)

test/test_blocksparsearrays.jl

@@ -130,9 +130,12 @@ arrayts = (Array, JLArray)
     @test_broken svd_compact(a)
   end
 
+  b = a[Block.(1:2), Block(2)]
+  @test b[Block(1)] == a[Block(1, 2)]
+  @test b[Block(2)] == a[Block(2, 2)]
+
   # Broken operations
   @test_broken exp(a)
-  @test_broken a[Block.(1:2), Block(2)]
 end
 
 @testset "BlockSparseArraysExt, EyeKronecker blocks (arraytype=$arrayt, eltype=$elt)" for arrayt in
@@ -174,19 +177,19 @@ end
   @test a[Block(2, 2)[(1:2) × (2:3), (:) × (2:3)]] ==
     a[Block(2, 2)][(1:2) × (2:3), (:) × (2:3)]
 
-  ## # Blockwise slicing, shows up in truncated block sparse matrix factorizations.
-  ## I1 = BlockIndexVector(Block(1), Base.Slice(Base.OneTo(2)) × [1])
-  ## I2 = BlockIndexVector(Block(2), Base.Slice(Base.OneTo(3)) × [1, 3])
-  ## I = [I1, I2]
-  ## b = a[I, I]
-  ## @test b[Block(1, 1)] == a[Block(1, 1)[(1:2) × [1], (1:2) × [1]]]
-  ## @test arg1(b[Block(1, 1)]) isa Eye
-  ## @test iszero(b[Block(2, 1)])
-  ## @test arg1(b[Block(2, 1)]) isa Eye
-  ## @test iszero(b[Block(1, 2)])
-  ## @test arg1(b[Block(1, 2)]) isa Eye
-  ## @test b[Block(2, 2)] == a[Block(2, 2)[(1:3) × [1, 3], (1:3) × [1, 3]]]
-  ## @test arg1(b[Block(2, 2)]) isa Eye
+  # Blockwise slicing, shows up in truncated block sparse matrix factorizations.
+  I1 = BlockIndexVector(Block(1), Base.Slice(Base.OneTo(2)) × [1])
+  I2 = BlockIndexVector(Block(2), Base.Slice(Base.OneTo(3)) × [1, 3])
+  I = [I1, I2]
+  b = a[I, I]
+  @test b[Block(1, 1)] == a[Block(1, 1)[(1:2) × [1], (1:2) × [1]]]
+  @test arg1(b[Block(1, 1)]) isa Eye
+  @test iszero(b[Block(2, 1)])
+  @test arg1(b[Block(2, 1)]) isa Eye
+  @test iszero(b[Block(1, 2)])
+  @test arg1(b[Block(1, 2)]) isa Eye
+  @test b[Block(2, 2)] == a[Block(2, 2)[(1:3) × [1, 3], (1:3) × [1, 3]]]
+  @test arg1(b[Block(2, 2)]) isa Eye
 
   # Slicing
   r = blockrange([2 × 2, 3 × 3])
@@ -272,7 +275,9 @@ end
   end
 
   # Broken operations
-  @test_broken a[Block.(1:2), Block(2)]
+  b = a[Block.(1:2), Block(2)]
+  @test b[Block(1)] == a[Block(1, 2)]
+  @test b[Block(2)] == a[Block(2, 2)]
 
   # svd_trunc
   dev = adapt(arrayt)