finch-tensor/finch-tensor-python

Slicing and indexing for `Tensor`

Closed this issue · 6 comments

mtsokol commented

Hi @willow-ahrens,

This is a tracking issue for implementing slicing and indexing functionality for Tensor class.

Right now I'm mapping objects that can be used for indexing from Python to Julia:

  • integers arr[1, 3] - this one only requires +1 operation to work.
  • slices (slice(1, 10, 2) or arr[1:10:2]) - when passing slices I get:
In [74]: arr_finch._obj.body[:,:,:]

JuliaError: ArgumentError: invalid index: <py slice(None, None, None)> of type Py
Stacktrace:
 [1] to_index(i::Py)
   @ Base ./indices.jl:300
 [2] to_index(A::Tensor{DenseLevel{Int64, DenseLevel{Int64, DenseLevel{Int64, ElementLevel{0, Int64, Int64, Vector{Int64}}}}}}, i::Py)
   @ Base ./indices.jl:277
 [3] _to_indices1(A::Tensor{DenseLevel{Int64, DenseLevel{Int64, DenseLevel{Int64, ElementLevel{0, Int64, Int64, Vector{Int64}}}}}}, inds::Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}, Base.OneTo{Int64}}, I1::Py)
   @ Base ./indices.jl:359
 [4] to_indices
   @ ./indices.jl:354 [inlined]
 [5] to_indices
   @ ./indices.jl:344 [inlined]
 [6] getindex(::Tensor{DenseLevel{Int64, DenseLevel{Int64, DenseLevel{Int64, ElementLevel{0, Int64, Int64, Vector{Int64}}}}}}, ::Py, ::Py, ::Py)
   @ Finch ~/.julia/packages/Finch/2wGtC/src/interface/index.jl:28
 [7] pyjlany_getitem(self::Tensor{DenseLevel{Int64, DenseLevel{Int64, DenseLevel{Int64, ElementLevel{0, Int64, Int64, Vector{Int64}}}}}}, k_::Py)

where Python slice is accepted by Julia arrays:

r = jl.rand(2,3,4)
r[:,:,:]
  • ellipsis ... - I don't think there's Julia counterpart of: arr[1, 2, ...] is arr[1, 2, :, :, :] when ndim(arr)==5, so I might construct it manually.
  • lists/arrays arr[[1,1,3]] - these also only require +1 operation on index arrays.
willow-ahrens commented

If the slice does not have a step of 1 I think we should try to treat it in Finch as if it's just some array.
If the slice has a step of 1 we do have a special case for that, but I think you would need to convert to the julia range type.
Could you convert the python slices to julia ranges?

willow-ahrens commented

I would indeed construct elipsis manually

mtsokol commented

Could you convert the python slices to julia ranges?

@willow-ahrens Oh, it looks that Julia's ranges can fully replace Python slices, thanks!:

In [110]: jl.size(arr_finch._obj.body[jl.range(start=1, step=2, stop=4),jl.range(length=2),jl.range(start=2, stop=3)])
Out[110]: (2, 2, 2)
willow-ahrens commented

I wonder why it wasn't captured in the automatic conversions listed here: https://juliapy.github.io/PythonCall.jl/stable/conversion-to-julia/

mtsokol commented

I think there's a conversion of Python's range to Julia's range, where both are iterable. I think that Python's slice object isn't converted.

willow-ahrens commented

Ah I see. There's a way to add conversions, though I'm not sure whether you think we should do that ourselves or through JuliaCall