[AutoDiff] [API] Fix bug in Tensor.reshaped VJP

[jekbradbury]

Introduced in #21740. Added a test.

[dan-zheng]

Nice catch, thanks!

[rxwei]

FYI: use expectEqual when you can, and the first argument is expected to be the expected value.

[dan-zheng]

If I recall, pure GPE can't handle expectEqual (because expectEqual opaquely calls Tensor.== and isn't inlined by GPE). That's also why test/TensorFlowRuntime/tensor.swift doesn't call Tensor.== directly.

Shall we use expectEqual anyways? I feel it especially makes sense for this test file.

[rxwei]

It can now, because #tfop has an ABI now.

[rxwei]

Nit:

• We generally don't use self. unless it's required.
• It'd be more efficient if you make it capture self.shapeTensor explicitly instead of capturing self. I'd write this as

{ [shape = shapeTensor] v in
  return v.reshaped(toShape: shape)
}

[dan-zheng]

FYI: there are definitely more places where we could use explicit captures instead of capturing self for efficiency. We should change those places at some point!

[rxwei]

Depends on whether self is practically a larger tensor that the result you get from operations on self :)

[rxwei]

@swift-ci please test tensorflow

[rxwei]

@swift-ci please test tensorflow

[rxwei]

@swift-ci please test tensorflow