Add support for CartesianIndex (#500)

jump-dev · Feb 10, 2023 · 0c1bb96 · 0c1bb96
1 parent d5e0208
commit 0c1bb96
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 0 deletions.
diff --git a/src/atoms/affine/index.jl b/src/atoms/affine/index.jl
@@ -133,4 +133,7 @@ getindex(x::AbstractExpr, cln_r::Colon, col) = getindex(x, 1:size(x)[1], col)
 # All columns for this row(s)
 getindex(x::AbstractExpr, row, cln_c::Colon) = getindex(x, row, 1:size(x)[2])
 
+# Cartesian Index
+getindex(x::AbstractExpr, c::CartesianIndex{N}) where {N} = x[Tuple(c)...]
+
 ## API Definition ends
diff --git a/test/test_utilities.jl b/test/test_utilities.jl
@@ -409,6 +409,17 @@ end
         @test size(y) == (2, 1)
     end
 
+    @testset "Cartesian index" begin
+        x = Variable(3, 2)
+        for ind in CartesianIndices(zeros(3, 2))
+            @test x[ind] === x[ind[1], ind[2]]
+        end
+        y = [1.0 2 3; 4 5 6] * x
+        for ind in CartesianIndices(zeros(2, 2))
+            @test y[ind] === y[ind[1], ind[2]]
+        end
+    end
+
     @testset "Parametric constants" begin
         z = Constant([1.0 0.0im; 0.0 1.0])
         @test z isa Constant{Matrix{Complex{Float64}}}