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):
- Indexing across multiple access at the same time.
tensor[2,6,-1]is preferable to
tensor[-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).
- Slicing across multiple axes without stride support. E.g.,
tensor[0..<5,3]. Note that partial ranges can also be supported in TensorFlow.
- Support for a
NewAxisliteral. This could potentially be done using
Nonecan be used in Python) or something like a singleton object (that’s how I represented this in TF Scala).
- 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.
- Support for an
Ellipsisliteral (this is
...in numpy). We could use a singleton object similar to
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
StrideThrough<T> object depending on the range it operates on. This would require the following additions to Swift:
- Add the new
- Expose the
- Ideally, add support for
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.