Advanced Array/Tensor Indexing

[eaplatanios (Anthony Platanios)]

Hi all,

While trying to use Swift for TensorFlow for a research project, I ran into the issue of missing advanced indexing functionality, like what's available for Python lists, Numpy arrays, or TensorFlow tensors in Python. The relevant discussion is here and the work-in-progress PR is here, but I will summarize things in this post, and also mention the currently unsolved issues in the end.

The following is a list of desired features, in terms of how often I believe each feature is used in Python (for TensorFlow tensors, all of the following features can be implemented using the TF strided slice op and so this is effectively a list of syntactic features):

1. Indexing across multiple access at the same time. tensor[2,6,-1] is preferable to tensor[2][6][-1] both in terms of syntax and implementation, as it is more efficient to do this in one operation (in terms of memory allocations and copies).
2. Slicing across multiple axes without stride support. E.g., tensor[0..<5,3]. Note that partial ranges can also be supported in TensorFlow.
3. Support for a NewAxis literal. This could potentially be done using nil (I believe None can be used in Python) or something like a singleton object (that’s how I represented this in TF Scala).
4. Support for strided slices. Basically, point 2 above with support for strided ranges. Not sure yet what a nice syntax for this would be in Swift.
5. Support for an Ellipsis literal (this is ... in numpy). We could use a singleton object similar to NewAxis for this.

Points 1, 2, 3, and 5 are all addressed in this PR.

Issue #1: Strided Ranges

Point 4 is not as easy to support currently, because there is no nice syntax for strided ranges. I wanted to see if anyone has good suggestions for how to go about this and also propose a way to do this.

What if we introduce a .. operator that has lower precedence than all the range operators (e.g., ..<, ..., etc), and that takes a range as first argument and a stride value as second. Example usage would look like 0..<5..2, where 2 is the stride. This would result in a StrideTo<T> or StrideThrough<T> object depending on the range it operates on. This would require the following additions to Swift:

• Add the new .. operator.
• Expose the lowerBound, upperBound, and stride properties of StrideTo<T> and StrideThrough<T>.
• Ideally, add support for PartialStrideTo<T>, PartialStrideThrough<T>, and PartialStrideFrom<T>.

What do people think about this?

Issue #2: Negative Indexing

This is lower priority, but Python and a lot of scientific computing libraries support negative indexing in arrays, lists, and tensors (e.g., -1 means 1 before the end). Could we introduce support for something similar in Swift? I understand this may be counterintuitive to people already used to Swift ranges not working that way, but there may be a nice solution to this that makes everyone happy.

Thanks,
Anthony

[Nevin]

It is not clear to me what you are asking for. Are you proposing changes to Swift language, or additions to the standard library, or modifications to TensorFlow?

@eaplatanios

• Indexing across multiple access at the same time. tensor[2,6,-1] is preferable to tensor[2][6][-1] both in terms of syntax and implementation, as it is more efficient to do this in one operation (in terms of memory allocations and copies).

This is easy to implement as an n-ary subscript on n-dimensional arrays. Any library which vends such a type can and should include such a subscript. I’m not sure what else you want—the language lets you make this already.

@eaplatanios

• Slicing across multiple axes without stride support. E.g., tensor[0..<5,3] . Note that partial ranges can also be supported in TensorFlow.

This can also be made as a subscript, although it would require 2ⁿ overloads to handle every combination of intermingled range expressions and single values.

Perhaps an approach here could be to conditionally conform ClosedRange to ExpressibleByIntegerLiteral, so the subscript would just take n parameters constrained to RangeExpression.

Alternatively, we could consider conforming Int to RangeExpression directly, but I imagine that would have unintended type-inference consequences.

@eaplatanios

• Support for a NewAxis literal. This could potentially be done using nil (I believe None can be used in Python) or something like a singleton object (that’s how I represented this in TF Scala).

I don’t know what this means. Could you briefly explain the desired behavior of NewAxis?

@eaplatanios

• Support for strided slices. Basically, point 2 above with support for strided ranges. Not sure yet what a nice syntax for this would be in Swift.

My memory is a little hazy, but I seem to recall previous discussions about this generally favored the syntax (a...b).by(2)

@eaplatanios

• Support for an Ellipsis literal (this is ... in numpy). We could use a singleton object similar to NewAxis for this.

Swift already has a standalone ellipsis operator, which is used for full-range slicing, eg, “let wholeSlice = someCollection[...]”. Do you have another use-case in mind, or are you talking about something different (if so please describe what you want)?

[scanon (Steve Canon)]

@eaplatanios

Indexing across multiple access at the same time. tensor[2,6,-1] is preferable to tensor[2][6][-1] both in terms of syntax and implementation, as it is more efficient to do this in one operation (in terms of memory allocations and copies).

You should focus on how the indexing is used. Efficiency is mostly a red-herring in this case--we should be able to make either of these perform fine; it may require some work on the optimizer, but we don't need to design libraries around that.

@eaplatanios

This is lower priority, but Python and a lot of scientific computing libraries support negative indexing in arrays, lists, and tensors (e.g., -1 means 1 before the end). Could we introduce support for something similar in Swift? I understand this may be counterintuitive to people already used to Swift ranges not working that way, but there may be a nice solution to this that makes everyone happy.

This makes me nervous that you're thinking of indexing as zero-based. Swift indexes are not, generally; specifically slices use the indexing of the collection they came from:

  1> let x: Array<Int> = [0,1,2,3]
x: [Int] = 4 values {
  [0] = 0
  [1] = 1
  [2] = 2
  [3] = 3
}
  2> let y = x[2...3]
y: ArraySlice<Int> = 2 values {
  [2] = 2
  [3] = 3
}
  3> y.startIndex
$R0: Int = 2

any "fancy indexing" proposal will need to account for this. I don't think that's a major blocker, but it does require some design, and it probably won't look exactly like it does in other languages that support these operations.

[eaplatanios (Anthony Platanios)]

Thanks for your responses! :)

@Nevin

This is easy to implement as an n-ary subscript on n-dimensional arrays. Any library which vends such a type can and should include such a subscript. I’m not sure what else you want—the language lets you make this already.

This, as well as most of the aforementioned points are already implemented in the linked PR by introducing a new protocol for valid indices.

Strided ranges are the main missing functionality.

@Nevin

I don’t know what this means. Could you briefly explain the desired behavior of NewAxis ?

NewAxis introduces a new dimension to multi-dimensional arrays/tensors. You can find more information on this here.

@Nevin

My memory is a little hazy, but I seem to recall previous discussions about this generally favored the syntax (a...b).by(2)

This is a bit too verbose in my opinion. For example, when doing a multi-dimensional slice you could get something like this:

tensor[(0..<3).by(2),(5...1).by(-1)]

as opposed to something like this:

tensor[0..<3..2, 5...1..(-1)]

@Nevin

Swift already has a standalone ellipsis operator, which is used for full-range slicing, eg, “ let wholeSlice = someCollection[...] ”. Do you have another use-case in mind, or are you talking about something different (if so please describe what you want)?

That's great to know! Thanks! We can actually use that operator directly I believe.

@scanon

You should focus on how the indexing is used . Efficiency is mostly a red-herring in this case--we should be able to make either of these perform fine; it may require some work on the optimizer, but we don't need to design libraries around that.

I agree, but the reason I mentioned this is because in this case (for TensorFlow), we already have a single op that is more efficient and can be built in this way, without making any changes to the Swift compiler.

@scanon

This makes me nervous that you're thinking of indexing as zero-based.

I'm not sure I understand what you mean. Is it that negative indices would obtain a different meaning, special-casing the beginning of a dimension to always be 0? Also, my understanding is that negative indices are currently not supported at all for Swift arrays. Is that right?

-Anthony

[scanon (Steve Canon)]

@eaplatanios

Is it that negative indices would obtain a different meaning, special-casing the beginning of a dimension to always be 0?

Suppose we supported negative indices on Array today, with the interpretation you have in mind:

let x = [0,1,2,3,4]
x[-1] // 4
let y = x[1...3] // shares indices with x.
y[-1] // Index out of bounds
y[-2] // 3

Also, negative indices never happen with Array, but there's no reason that another type conforming to Collection couldn't use them (Indices only need to be Comparable). Tensors don't necessarily need to conform to Collection, but they shouldn't create surprises for people who expect them to behave roughly like other indexable Swift types.

Swift does have an extremely verbose way to spell this today:

y[y.index(y.endIndex, offsetBy: -1, limitedBy: y.startIndex)!]

I think it would be great to provide a simpler way of expressing what is fundamentally integer arithmetic, but it will necessarily be a bit more complex than just a negative subscript. (CC @Michael_Ilseman who has sketched some ideas in this direction before).

[eaplatanios (Anthony Platanios)]

@Nevin

Swift already has a standalone ellipsis operator, which is used for full-range slicing, eg, “ let wholeSlice = someCollection[...] ”. Do you have another use-case in mind, or are you talking about something different (if so please describe what you want)?

It looks like I can't actually use the existing ... operator in the current implementation because it is a function type, (UnboundedRange_) -> (), and I cannot define protocol conformances for function types. Why is it a function type though? Is it because it needs to be defined as an operator? Couldn't it be simply defined as a top-level global variable?

[eaplatanios (Anthony Platanios)]

@scanon

I think it would be great to provide a simpler way of expressing what is fundamentally integer arithmetic, but it will necessarily be a bit more complex than just a negative subscript. (CC @Michael_Ilseman who has sketched some ideas in this direction before).

I agree. I think this would be great. A starting point may be a less verbose syntax for referencing the start and end indices of an array/tensor, and then being able to use normal arithmetic operators on that.