[TF] Added support for advanced indexing and slicing #23684

eaplatanios · 2019-03-30T04:53:01Z

This adds support for advanced indexing and slicing to tensors. I plan to add support for setters and gradients soon, but first I wanted to see if this initial design looks good to the Swift for TF team.

I also added a new test and made sure all previous tests using tensor subscripts still pass, using the new subscript implementation.

Some comments/questions:

When using a subscript setter in code we get a crash, as in TF-178. I tried looking into this a bit and it seems to me that the type checker is crashing when trying to diagnose a subscript error in this case. This is probably because, in the failure example, we're trying to use the subscript setter, but no such setter exists, and there is bug in how this error is diagnosed. There does not seem to be an issue when using the getter. I'm not sure how to check if this has been fixed upstream as the merging upstream changes may be nontrivial (although I haven't tried doing so yet).
The overload subscript functions were required due not having support for variadic generics. Do you know of a better way to do this? Also, what is the convention for line breaks in the generic types list?
Regarding gradients, is it possible to use @differentiating to directly provided the gradient of Raw.stridedSlice? Otherwise, I can provide a gradient directly for the subscript function.
Swift ranges do not support negative indices and so something like 0 ..< -1 cannot be used, even if only as syntactic sugar. Do we need that precondition for Swift ranges? In principle, the ranges could be allowed to be empty, from a Swift point of view, and have this be used simply as syntactic sugar. This can also be avoided if we define our own ..<, ..., etc. operators that support negative indices and are only used as syntactic sugar for advanced indexing and slicing.
Strides are not currently supported, but I can post a question on Swift Evolution about this tomorrow. One very direct but S4TF-specific solution would be to provide our own ..<, ..., etc. operators as discussed above, that also allow for a .. operator for strided ranges.
Is there a way to also support ranges over Int32, Tensor<Int32>, and other types? I get an error whenever I try to do that saying I can only define one extension for each protocol conformance, irrespective of the conditional constraint on the type.
We should probably start organizing the ops in multiple files instead of putting everything in Ops.swift. How about making an Ops directory in TensorFlow and have separate files like Manipulation.swift, Math.swift, etc.

stdlib/public/TensorFlow/Ops.swift

rxwei · 2019-03-30T06:15:33Z

When using a subscript setter in code we get a crash, as in TF-178. I tried looking into this a bit and it seems to me that the type checker is crashing when trying to diagnose a subscript error in this case. This is probably because, in the failure example, we're trying to use the subscript setter, but no such setter exists, and there is bug in how this error is diagnosed. There does not seem to be an issue when using the getter. I'm not sure how to check if this has been fixed upstream as the merging upstream changes may be nontrivial (although I haven't tried doing so yet).

Folks are still working on merging things from master. I'll take a little longer.

Have you tried adding a setter?

The overload subscript functions were required due not having support for variadic generics. Do you know of a better way to do this? Also, what is the convention for line breaks in the generic types list?

Before we have variadic generics, we can use existentials.

subscript(_ indices: TensorSliceIndexProtocol...)

The Google Swift Style Guide is slightly inconsistent with out local convention: https://google.github.io/swift/#line-wrapping. But you can follow its line wrapping guide.

Regarding gradients, is it possible to use @differentiating to directly provided the gradient of Raw.stridedSlice? Otherwise, I can provide a gradient directly for the subscript function.

Not yet. The feature is called "retroactive derivative registration" and will take some time to engineer.

Swift ranges do not support negative indices and so something like 0 ..< -1 cannot be used, even if only as syntactic sugar. Do we need that precondition for Swift ranges? In principle, the ranges could be allowed to be empty, from a Swift point of view, and have this be used simply as syntactic sugar. This can also be avoided if we define our own ..<, ..., etc. operators that support negative indices and are only used as syntactic sugar for advanced indexing and slicing.

We should stick with Swift's range formation operators and semantics for now. The precondition that upperBound cannot be -1 (or any number smaller than lowerBound) is in Swift programmers' mental model IMO. Let's delay the discussion about negative indices after this PR.

Strides are not currently supported, but I can post a question on Swift Evolution about this tomorrow. One very direct but S4TF-specific solution would be to provide our own ..<, ..., etc. operators as discussed above, that also allow for a .. operator for strided ranges.

I don't think providing our own ..< and ... solves the problem here. I believe we need an operator (of lower precedence than RangeFormationPrecedence) that takes a Range/ClosedRange on the left hand side and takes an integer on the right hand side, forming a StridedRange. The only new things necessary are StridedRange and that operator.

Is there a way to also support ranges over Int32, Tensor<Int32>, and other types? I get an error whenever I try to do that saying I can only define one extension for each protocol conformance, irrespective of the conditional constraint on the type.

For now, I'd recommend against supporting Tensor<Int32> indices since it can be more error-prone. Int32 is the index type for Tensor APIs, and there is no precedent of supporting multiple index types in the Swift standard library. Let's stick with Int32 for now.

We should probably start organizing the ops in multiple files instead of putting everything in Ops.swift. How about making an Ops directory in TensorFlow and have separate files like Manipulation.swift, Math.swift, etc.

Actually, our plan is to gradually move everything to tensorflow/swift-apis.

eaplatanios · 2019-03-30T17:32:22Z

Have you tried adding a setter?

I'll add this today. I can use TensorScatterUpdate using #tfop, but I was wondering why it is not included in RawOpsGenerated.swift. Is it skipped because something is not supported yet?

Not yet. The feature is called "retroactive derivative registration" and will take some time to engineer.

Ok, I can use @differentiable in that case for now. One question related to that is how to I specify VJP separately for a subscript getter and setter using @differentiable?

We should stick with Swift's range formation operators and semantics for now. The precondition that upperBound cannot be -1 (or any number smaller than lowerBound) is in Swift programmers' mental model IMO. Let's delay the discussion about negative indices after this PR.

Sounds good.

I don't think providing our own ..< and ... solves the problem here. I believe we need an operator (of lower precedence than RangeFormationPrecedence) that takes a Range/ClosedRange on the left hand side and takes an integer on the right hand side, forming a StridedRange. The only new things necessary are StridedRange and that operator.

Sorry, that's what I meant with the .. operator for strided ranges. I just mixed the two comments in one making it confusing. Also, we don't necessarily need StridedRange given that Swift already has StrideTo<T> and StrideThrough<T> types. We would just need to expose their lower/upper bounds and strides as public properties. I'll include that in the Swift Evolution post.

For now, I'd recommend against supporting Tensor<Int32> indices since it can be more error-prone. Int32 is the index type for Tensor APIs, and there is no precedent of supporting multiple index types in the Swift standard library. Let's stick with Int32 for now.

Sounds good, although, if we want to avoid using gather (I noticed it's more expensive due to memory allocations based on the existing comments), we should probably add support for that soon as it's useful for embedding table lookups.

Also, the strided slice gradient is a dense tensor, so we can support it fine for now, but an advantage of gather is that, once we support sparse tensor gradients, we can compute its gradient as something equivalent to TensorIndexedSlices which would be more efficient in some cases.

Actually, our plan is to gradually move everything to tensorflow/swift-apis.

That sounds great actually! I wonder if we should be making these new additions for ops directly in that repo in that case. Let's keep this PR here for now, but maybe next time I work on adding support for new ops, I may do it there directly.

eaplatanios · 2019-03-30T17:36:13Z

Also, I just realized that the ScatterUpdate op in the raw ops module is generated wrongly here. This should instead be the signature of TensorScatterUpdate. For ScatterUpdate the ref tensor type should not be T, but rather the reference type for T, which is not supported in S4TF. This is because this op is meant to do updates on reference variables.

rxwei · 2019-03-30T22:20:34Z

Sorry, that's what I meant with the .. operator for strided ranges. I just mixed the two comments in one making it confusing. Also, we don't necessarily need StridedRange given that Swift already has StrideTo<T> and StrideThrough<T> types. We would just need to expose their lower/upper bounds and strides as public properties. I'll include that in the Swift Evolution post.

Great, we are on the same page. Before StrideTo and StrideThrough go through Swift Evolution, adding a new type is fine.

That sounds great actually! I wonder if we should be making these new additions for ops directly in that repo in that case. Let's keep this PR here for now, but maybe next time I work on adding support for new ops, I may do it there directly.

Feel free to move APIs there in future PRs.

Also, I just realized that the ScatterUpdate op in the raw ops module is generated wrongly here. This should instead be the signature of TensorScatterUpdate. For ScatterUpdate the ref tensor type should not be T, but rather the reference type for T, which is not supported in S4TF. This is because this op is meant to do updates on reference variables.

The raw op generation script should be changed to ignore ops that take ref arguments. As to TensorScatterUpdate, I'm not sure whether Raw.tensorScatterUpdate is available in the raw module yet. Maybe we need to regenerate bindings.

eaplatanios · 2019-03-30T22:25:26Z

I looked into the setter a bit and I realized that TensorScatterUpdate will not be efficient for this operation. We should directly add support for StridedSliceAssign, but for tensors (e.g., we can call it TensorStridedSliceUpdate). This should be easily done by taking the implementation of the existing StridedSliceAssign kernels and adding support for tensor arguments by following a similar approach to the TensorScatterUpdate kernel implementation. Does that sound reasonable?

eaplatanios · 2019-03-30T22:26:02Z

The raw op generation script should be changed to ignore ops that take ref arguments. As to TensorScatterUpdate, I'm not sure whether Raw.tensorScatterUpdate is available in the raw module yet. Maybe we need to regenerate bindings.

I agree. Also, Raw.tensorScatterUpdate is not available in the raw module yet.

eaplatanios · 2019-03-30T22:58:16Z

I added support for a new TensorStridedSliceUpdate op here. Haven't tested this yet, but if it works, would be able to merge this into the main TensorFlow repo and use it for the setter? Otherwise, is there currently a way to add support for new ops and kernel in S4TF?

rxwei · 2019-03-30T23:11:54Z

Haven't tested this yet, but if it works, would be able to merge this into the main TensorFlow repo and use it for the setter?

You can start a PR and go through reviews in core TF with folks like @alextp. This looks right, but I don't own the kernel code in core TensorFlow to say this is the right direction or not.

Otherwise, is there currently a way to add support for new ops and kernel in S4TF?

There's no way yet. Adding the kernel in core TF will be the first step.

eaplatanios · 2019-03-31T13:56:43Z

I'm following through with Alex on the core TF PR. In the meantime, how can I add separate VJP registrations for the getter and the setter of a subscript?

rxwei · 2019-03-31T18:50:53Z

Setters cannot have a VJP yet.

eaplatanios · 2019-03-31T18:52:03Z

Ok, so that means I can just use @differentiating on top of subscript itself, right?

rxwei · 2019-03-31T19:28:38Z

Use @differentiable(vjp: ... where ...) on top of the subscript.

eaplatanios · 2019-04-09T17:08:15Z

@rxwei I updated the TensorFlow checkout commit to include the strided slice assign op that was merged, but when I try to compile, I get a linking error saying that dyld: Library not loaded: @rpath/libtensorflow.so.1. Not sure if the dynamic libraries were versioned before, but if they were not, could this be due to something missing from the build script in S4TF?

rxwei · 2019-04-09T17:20:11Z

Could you try a bazel build in ../tensorflow?

eaplatanios · 2019-04-09T17:50:42Z

That succeeds. I can also see the compiled binaries in `../tensorflow/bazel-bin/tensorflow`.

…

On Tue, Apr 9, 2019 at 1:21 PM Richard Wei ***@***.***> wrote: Could you try a bazel build in ../tensorflow? — You are receiving this because you authored the thread. Reply to this email directly, view it on GitHub <#23684 (comment)>, or mute the thread <https://github.com/notifications/unsubscribe-auth/ABPCXL0qI4wN9fICdRMGY9bvF9lK04sCks5vfMwngaJpZM4cTioG> .

rxwei · 2019-04-09T17:52:41Z

You can make a renamed copy of the lib and try that.

eaplatanios · 2019-04-09T17:56:28Z

But where should I place it? ../tensorflow/bazel-bin/tensorflow already contains libtensorflow.so, libtensorflow.so.1 and libtensorflow.so.13.1, where the first two are symbolic links pointing to the last one.

eaplatanios · 2019-04-09T20:43:23Z

@rxwei I think this may be related to this PR and not the swift build configuration.

eaplatanios · 2019-04-16T23:48:08Z

Yes I agree. By the way, what's the difference between a nullary operator and a global variable? I tried to start an evolution topic that included strides, but I guess it may have been too generic and not focused on a single subject.

By the way, I just pushed the requested changes. If these look ok, I can go ahead and merge upstream changes including the Int32 -> Int change.

rxwei · 2019-04-16T23:50:23Z

Operators and identifiers are separate sets of characters. That's the main difference. Operators are functions, so I guess a nullary operator would have to represent a nullary function. Even if we had nullary operators, it still won't solve our problem. We need structural types (including function types) to be able to conform to protocols.

rxwei

By the way, I just pushed the requested changes. If these look ok, I can go ahead and merge upstream changes including the Int32 -> Int change.

LGTM! One suggestion is adding tests for the newly added .. operators.

stdlib/public/TensorFlow/Ops.swift

eaplatanios · 2019-04-17T16:00:40Z

@rxwei This should be ready now. I have also added a couple of tests for the new stride operator.

rxwei · 2019-04-17T18:06:31Z