Restructured the operators source files and added support for multiple new operators.

[eaplatanios]

This PR does some restructuring of the files where ops are defined and also adds the following:

• Constructors: Tensor(stacking:alongAxis:), Tensor(concatenating:alongAxis:), and Tensor(tiling:multiples:), along with their VJPs. Also adds Tensor(zerosLike:), Tensor(onesLike:), and Tensor(rangeFrom:to:stride:) with tensor-valued arguments.
• Ops: Adds Tensor.unstack(), Tensor.split(numSplits:alongAxis:), Tensor.split(sizes:alongAxis:), Tensor.gathering(atIndices:alongAxis:),sigmoid(_:), logSigmoid(_:), and softplus(_:) along with their VJPs. Also adds Tensor.batchGathering(atIndices:alongAxis:numBatchDims:), Tensor.gathering(where:alongAxis:), and Tensor.nonZeroIndices().
• Protocols: Adds an Optimizable protocol that makes the optimizers more generally applicable, other than just for layers.

Current issues:

• Cannot backprop through precondition which prevents me from making some of the new functions differentiable.

This is a work-in-progress PR that also depends on the advanced indexing PR, but I want to start a discussion and gather feedback about the source files restructuring and also some missing op support. I actually have the following main questions:

- Is the restructuring of the operators source files acceptable? I believe it will be hard to have all operators in a single source file.
- I currently cannot add support for gradients of stacked and concatenated with multiple tensors because I see that #tfop does not support ops that output lists of tensors. Are there any updates on this or any current workarounds to return arrays of tensors?
- Is it ok to merge PRs that add support for new ops but not VJPs yet? I plan to add those to incrementally.

This PR so far adds support for the stacked, concatenated, gathered, batchGathered, masked, and nonZeroIndices ops, as well as for a new tensor initializer from ranges defined over tensors, and new stacking, concatenating, and tiling tensor initializers.

[eaplatanios]

@rxwei What should I do about the #tfop limitation mentioned above?

[rxwei]

#tfop does support input lists and output lists. It supports anything that conforms to TensorArrayProtocol. Just use an array as arguments and/or results.

[eaplatanios]

#tfop does support input lists and output lists. It supports anything that conforms to TensorArrayProtocol. Just use an array as arguments and/or results.

In that case, why does the Raw module generation script say it doesn't? Is it just that script that does not support them?

[rxwei]

In that case, why does the Raw module generation script say it doesn't? Is it just that script that does not support them?

For historical reasons. GPE doesn't support variable-sized inputs/outputs, and swift-bindings started to exist when only GPE existed. We never updated the raw op generation script to generate array types, after eager execution became possible.

[eaplatanios]

In that case, I made a PR to add support for them in the Raw module generation script.

[eaplatanios]

I also addressed your comments, but will add the missing VJPs tomorrow because I have to finish up some other stuff tonight.

[rxwei]

Please move the "explicit paddings" fix to a separate PR so that it can get merged before all the dependencies of this PR can. Thanks!

[eaplatanios]

This compiles and passes all tests on my machine, when I follow the following process:

• Copy this over to the tensorflow-swift-apis directory (same as what the update-checkouts command would do).
• Compile apple/swift, run all tests, and build toolchain.

This process all works fine.

The problem is when I try to compile this on its own using the built toolchain from above. That toolchain already contains all symbols from swift-apis in the stdlib, as it builds swift-apis together with the stdlib. This means that if I try to compile this with that toolchain I get many errors about ambiguous functions.

[rxwei]

Yeah, I understand the apple/swift path compiles and passes tests. To ensure swift-apis can be developed individually, we'll also need to make sure it works with Swift Package Manager.

Earlier there were already quite a few operators and layers that were both defined in here and distributed as part of the toolchain. There were no ambiguity errors because overriding a symbol in a different module is okay. There may be some subtleties in the newly migrated APIs. It would be useful to see what exactly those ambiguities are.

[rxwei]

Let me know if you have any question via email or a new PR. I'm more than happy to help with those issues!

[eaplatanios]

I see, yes that makes sense. So, the output when compiling swift-apis looks as follows:

• Multiple warnings about protocol conformances that were declared in stdlib. This is expected since the "namescope" does not matter in this case.
• Multiple errors about the @_semantics attribute I added in the last PR. These look like: error: '@_semantics' attribute cannot be applied to this declaration.
• Error in PythonConversion.swift: use of unresolved identifier 'debugLog'.
• Error that I can't figure out:

/Users/eaplatanios/Development/GitHub/swift-apis/Sources/DeepLearning/Operators/Basic.swift:342:48: error: ambiguous use of 'subscript(_:)'
        let splits = Tensor<Int32>([shapeTensor[idx], other.shapeTensor[idx]])
                                               ^
TensorFlow.Tensor<τ_0_0>:2:23: note: found this candidate
    @inlinable public subscript(ranges: TensorRangeExpression...) -> TensorFlow.Tensor<Scalar> { get set }
                      ^
/Users/eaplatanios/Development/GitHub/swift-apis/Sources/DeepLearning/Operators/Basic.swift:633:17: note: found this candidate
    fileprivate subscript(_ ranges: TensorRangeExpression...) -> Tensor {
                ^

[eaplatanios]

That's all actually so it should be easy to resolve.

[rxwei]

I missed this newly added operator. Swift APIs usually don't use the "num" prefix to denote count, but use a Count suffix instead. In this case, a single count may also be fine. You could call this split(count:alongAxis:) or follow the standard library precedent: split(maxSplits:omittingEmptySubsequences:whereSeparator:).

[rxwei]

Functions that do not have side effects should be named like nouns (see Swift API Design Guidelines - Naming - Strive For Fluent Usage). Since this method doesn't have side effects, it a gerund (unstacking(alongAxis:)) would be a better name.

[rxwei]

• Multiple errors about the @_semantics attribute I added in the last PR. These look like: error: '@_semantics' attribute cannot be applied to this declaration.

I think that @_semantics is only available on explicit property getters (get {}).

• Error in PythonConversion.swift: use of unresolved identifier 'debugLog'.

Feel free to comment out those debugLogs for now!

• Error that I can't figure out:

/Users/eaplatanios/Development/GitHub/swift-apis/Sources/DeepLearning/Operators/Basic.swift:342:48: error: ambiguous use of 'subscript(_:)'
        let splits = Tensor<Int32>([shapeTensor[idx], other.shapeTensor[idx]])
                                               ^
TensorFlow.Tensor<τ_0_0>:2:23: note: found this candidate
    @inlinable public subscript(ranges: TensorRangeExpression...) -> TensorFlow.Tensor<Scalar> { get set }
                      ^
/Users/eaplatanios/Development/GitHub/swift-apis/Sources/DeepLearning/Operators/Basic.swift:633:17: note: found this candidate
    fileprivate subscript(_ ranges: TensorRangeExpression...) -> Tensor {
                ^

I'm not sure either. Why is the one on 633 fileprivate in the error message?

[eaplatanios]

I think that @_semantics is only available on explicit property getters (get {}).

That fixes it!

Feel free to comment out those debugLogs for now!

Done.

I'm not sure either. Why is the one on 633 fileprivate in the error message?

This was just something I was trying. Could this error be because the subscript function is defined in a tensor extension and that may also cause a conflict when redefined across different modules? TensorFlow is imported in that file and so there is already a Tensor.subscript with the same signature defined in scope.

[rxwei]

I don't have a good solution to that. How about defining indexing logic in stdlib for now?

[eaplatanios]

That's ok for now but I also believe it is a more important issue that'll likely pop up again with other extensions.

[eaplatanios]

Actually I believe this might happen a lot in the future, given that a lot of the added functionality is in the form of extensions to Tensor. Although, I am curious as to why we're only observing this for subscript and not for other functions that are also implemented in extensions of Tensor.

[rxwei]

Yeah, I think we should definitely get it resolved sooner rather than later. The Foundation model would be good to follow. Other people on the team have talked about moving TensorFlow out of the compiler repo as well (@pschuh @jekbradbury).

[eaplatanios]

I would be happy to help with that because even the changes introduced in this PR seem hard to properly handle (in that swift-apis cannot easily be compiled independently using S4TF). And yes, I would lean in favor of moving TensorFlow out of the compiler repo if that is an option. Could you please actually describe the Foundation model at a high level?

[eaplatanios]

Most importantly actually, I'm curious about how the deal with the ambiguous symbols issue that we are running into.

[rxwei]

It's really awesome to hear you are interested in helping out!

Foundation sources are located in apple/swift-corelibs-foundation, and it's being pulled into the build directory by update-checkout. When building using build-toolchain, it is being included as one of the core libraries in toolchains so that users can import Foundation.

Foundation is a little different in that it's not a SwiftPM package, but I suspect that's not the difficult part.

[rxwei]

Most importantly actually, I'm curious about how the deal with the ambiguous symbols issue that we are running into.

Are you referring to the subscript errors? I'd suggest pulling them back to stdlib for now.

[eaplatanios]

I've started using S4TF for some new research projects and would generally love to make the process of improving it and extending it smoother. I love the work you guys have put in and all that lies ahead! :)

Regarding foundation, that sounds similar to what's happening now with swift-apis. I believe the main challenge (irrespective of using SwiftPM) is in compiling swift-apis as a standalone library using the S4TF compiler. I believe this would be much easier if we move all of the TensorFlow library outside of the compiler repo and into swift-apis. I would be willing to give this a shot, although I'm not sure if there is anything internal being used in the TensorFlow directory in stdlib that won't compile if I move it out. I could give it a shot as part of #109 and see if things work out.

[rxwei]

I've started using S4TF for some new research projects and would generally love to make the process of improving it and extending it smoother. I love the work you guys have put in and all that lies ahead! :)

That's fantastic. We really really appreciate your effort in driving important changes as well as pushing the limits of the incomplete Swift for TensorFlow.

I believe the main challenge (irrespective of using SwiftPM) is in compiling swift-apis as a standalone library using the S4TF compiler.

Right. This is why we probably shouldn't attempt this right now. In the TensorFlow module there is legacy code w.r.t. GPE, and also new experimental code built by @bgogul and @pschuh et al.

One option is to keep TensorHandle (as well as other important handle types) in the compiler repo as TensorFlowCore, and all other things in a separate TensorFlow repo.

I feel that tensor APIs being in the compiler repo really starts to limit our development. So it's probably good time to get the discussion started. How about setting up a meeting with the team to discuss this? Are you able to video chat some time this or next week?

[vguerra]

hello @eaplatanios. Quick question on this PR. Description claims the addition of initializers: Tensor(zerosLike:) and Tensor(onesLike:) but I don't see them, or maybe I am missing something (?):/.

Same goes for protocol Optimizable.

[rxwei]

Not sure about the new initializers. The ‘Optimizable’ protocol has been reverted, addressing this review feedback.

[vguerra]

@rxwei thanks.. missed those comments.

[eaplatanios]

I feel that tensor APIs being in the compiler repo really starts to limit our development. So it's probably good time to get the discussion started. How about setting up a meeting with the team to discuss this? Are you able to video chat some time this or next week?

Yes, definitely. That would be great! How about sometime on Wednesday?

Also, I like the TensorFlowCore suggestion and that would resolve the ambiguous symbols issues too I believe.

[eaplatanios]

hello @eaplatanios. Quick question on this PR. Description claims the addition of initializers: Tensor(zerosLike:) and Tensor(onesLike:) but I don't see them, or maybe I am missing something (?):/.

Same goes for protocol Optimizable.

Sorry, yes we decided to remove the new ops from this PR to make it about moving stuff over from stdlib first. I have all the new ops implemented but will add them in a separate PR, after we sort out #109 .

[vguerra]

Sorry, yes we decided to remove the new ops from this PR to make it about moving stuff over from stdlib first. I have all the new ops implemented but will add them in a separate PR, after we sort out #109 .

Got it! .. thanks!

[rxwei]

Yes, definitely. That would be great! How about sometime on Wednesday?

Reached out over email regarding the meeting!