Generating function for multi-dimensional array indexing

General Usage
, ,
1

Hey all,

I’m trying to generate some functions, and while they work, I don’t fully understand the behavior coming out. Also, there may be some function in Base which does this, I just can’t find it. Illumination of either would be very appreciated :slight_smile:

I want to generate a function slice to get a slice 2D slice of a 3D array. slice should take two args: an int for which dimension should be ‘flat’, and the value for the index to take along that dimension.
Eg: slice(A, 2, 3) should give me A[:, 3, :], slice(A, 1, 2) should give A[2, :, :].
Am I reinventing a wheel here?

From the docs, @generated ‘generate[s] specialized code depending on the types of their arguments’. I expected a method for each set of types passed in, but we’re seeing just one. How is the code specialized then? I want to ask ‘where’ the specialized code is (it if that makes sense).

Here is my MWE:

access_expression(F, i) = Expr(vcat([:ref, F], circshift([:x_i, :, :], i - 1))...)
dump(access_expression(:A, 2))

@generated function slice(A, i :: Val{T}, x_i) where T
    # we want @generated to dispatch on dimension
    slice_dimension = T
    quote
        $(access_expression(:A, slice_dimension))
    end
end

A = rand(Int16, 1, 2, 3)
println("A = \n$A")
@assert size(slice(A, Val(1), 1)) == (2, 3)
println("slice 1 = \n $(slice(A, Val(1), 1))")
@assert size(slice(A, Val(2), 1)) == (1, 3)
println("slice 2 = \n $(slice(A, Val(2), 1))")
@assert size(slice(A, Val(3), 1)) == (1, 2)
println("slice 3 = \n $(slice(A, Val(3), 1))")

which gives us the following output:

  head: Symbol ref
  args: Array{Any}((4,))
    1: Symbol A
    2: Colon() (function of type Colon)
    3: Symbol x_i
    4: Colon() (function of type Colon)
A = 
Int16[-9408 10291]

Int16[-28752 -4734]

Int16[28341 -30523]
slice 1 =
 Int16[-9408 -28752 28341; 10291 -4734 -30523]
slice 2 =
 Int16[-9408 -28752 28341]
slice 3 =
 Int16[-9408 10291]

julia> slice
slice (generic function with 1 method)

We can see from dump that the expression produced is simple, and dimension specific. The slicing works. Yet there’s only one method!

2

This sounds like selectdim.

2 Likes
3

I believe selectdim should already do what you want.

As for your question about the number of methods, the important thing to note is that specialization is independent of the number of methods of a function. When you write f(x) = 2x and then call f(1) and f(1.0), Julia specializes f for Int and Float64 inputs, but those are both specializations of the single method f(x) = 2x.

4 Likes