[TF] Fix gradients in sum(squeezinAxes:) and mean(squeezinAxes:)

[sgugger]

sum(squeezinAxes:) and mean(squeezinAxes:) were throwing an error during the backward pass because the gradients weren't unsqueezed before being broadcast, this PR fixes it.

Note that this could be refactored nicely if we had a function that took a list of ints for expandingShape.
Second note: I may be wrong, but it seems like _vjpMean(squeezingAxes axes: [Int]) is never used and only the Tensor version is.

[dan-zheng]

Reviewing now.

[dan-zheng]

It would be nice to add tests and an expandingShape(at shapeIndices: [Int]) API, but those needn't block progress.

[dan-zheng]

@swift-ci Please test tensorflow

[sgugger]

I can work on an expandingShape(at shapeIndices: [Int]). Not sure where to add tests, but I can look.

[dan-zheng]

I can work on an expandingShape(at shapeIndices: [Int]).

That would be great, if you have time!

Not sure where to add tests, but I can look.

Please add tests to test/TensorFlowRuntime/tensor_autodiff_runtime.swift.

[rxwei]

Thanks!

[rxwei]

I can work on an expandingShape(at shapeIndices: [Int]). Not sure where to add tests, but I can look.

That'd be great. You can check it into tensorflow/swift-apis if that's easier!

[dan-zheng]

@swift-ci Please test tensorflow

[sgugger]

There are actually commented tests for all squeezing things in the file you mentioned Dan. Is there any reason for that?

[dan-zheng]

There are actually commented tests for all squeezing things in the file you mentioned Dan. Is there any reason for that?

Great catch! squeezingAxes: derivative tests were commented in #22112, exactly because the behavior was not correct/well-defined.

This PR fixes those concerns. Are you in a position to uncomment the tests and run tests locally? If not, I can do so. Please let me know.

[sgugger]

I can update with the uncommented tests, just not sure how to run them locally.

[dan-zheng]

I can update with the uncommented tests, just not sure how to run them locally.

Okay. Please update the tests, and I will verify that they pass locally, then trigger CI. Thanks!

[sgugger]

Done, when you have a bit of time, please tell me how you can quickly test one of those files so that I don't bother you next time ;)

[dan-zheng]

Done, when you have a bit of time, please tell me how you can quickly test one of those files so that I don't bother you next time ;)

Thanks! Running tests requires a local build of the Swift compiler - if you don't have a local build, quickly testing is difficult.

If you do have a local build, you can run tests using lit.py (the LLVM Integrated Tester). Here's a detailed guide from the Swift forums.

To run this particular test, I run the following (you may need to replace macosx-x86_64):

cd swift-dev # cd into directory containing `build` and `swift`
llvm/utils/lit/lit.py -vvv --param swift_test_mode=optimize build/Ninja-ReleaseAssert+stdlib-Release/swift-macosx-x86_64/test-macosx-x86_64/TensorFlowRuntime/tensor_autodiff_runtime.swift

[dan-zheng]

Verified that tests pass locally. Triggering CI.

[dan-zheng]

@swift-ci Please test tensorflow

[dan-zheng]

@swift-ci Please test tensorflow linux

[rxwei]

Thanks all!