Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 5 commits into from
Apr 19, 2019
Merged

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
3ac59cd
[TF] Fix gradients in sum(squeezinAxes:) and mean(squeezinAxes:)
sgugger Apr 19, 2019
fb565e0
Remove unused `_vjpMean` function.
dan-zheng Apr 19, 2019
b2fa44a
Update Gradients.swift
dan-zheng Apr 19, 2019
dfbd990
Add test
sgugger Apr 19, 2019
65752cf
Minor edit for consistency.
dan-zheng Apr 19, 2019
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 8 additions & 11 deletions stdlib/public/TensorFlow/Gradients.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -579,7 +579,11 @@ extension Tensor where Scalar : TensorFlowFloatingPoint {
squeezingAxes axes: Tensor<Int32>
) -> (Tensor, (Tensor) -> Tensor) {
let value = sum(squeezingAxes: axes)
return (value, { [shape = shapeTensor] in $0.broadcast(toShape: shape) })
return (value, { [shape = shapeTensor] in
var res = $0
for i in axes.array.scalars { res = res.expandingShape(at: Int(i)) }
return res.broadcast(toShape: shape)
})
}

@inlinable
Expand All @@ -591,23 +595,16 @@ extension Tensor where Scalar : TensorFlowFloatingPoint {
})
}

@inlinable
func _vjpMean(squeezingAxes axes: [Int]) -> (Tensor, (Tensor) -> Tensor) {
let value = mean(squeezingAxes: axes)
return (value, { [shape = shapeTensor,
count = axes.map { shape[$0] }.reduce(1, *)] in
$0.broadcast(toShape: shape) / Tensor(Scalar(count))
})
}

@inlinable
func _vjpMean(
squeezingAxes axes: Tensor<Int32>
) -> (Tensor, (Tensor) -> Tensor) {
let value = mean(squeezingAxes: axes)
let count = Raw.gather(params: shapeTensor, indices: axes).product()
return (value, { [shape = shapeTensor] in
$0.broadcast(toShape: shape) / Tensor(count)
var res = $0
for i in axes.array.scalars { res = res.expandingShape(at: Int(i)) }
return res.broadcast(toShape: shape) / Tensor(count)
})
}
}
Expand Down
13 changes: 7 additions & 6 deletions test/TensorFlowRuntime/tensor_autodiff_runtime.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -98,38 +98,39 @@ TensorADTests.testAllBackends("Abs") {
TensorADTests.testAllBackends("sum") {
let input = Tensor<Float>(repeating: 42, shape: [2, 2])
let sumPullbackScalar = pullback(at: input) { (a: Tensor<Float>) in a.sum() }
let sumPullbackSqueezingAxes = pullback(at: input) { (a: Tensor<Float>) in a.sum(squeezingAxes: 0, 1) }
let sumPullbackAlongAxes = pullback(at: input) { (a: Tensor<Float>) in a.sum(alongAxes: 0, 1) }

let expected = Tensor<Float>(ones: [2, 2])
expectEqual(expected, sumPullbackScalar(Tensor(1)))
// expectEqual(expected, sumPullbackSqueezingAxes(Tensor(1)))
expectEqual(expected, sumPullbackSqueezingAxes(Tensor(1)))
expectEqual(expected, sumPullbackAlongAxes(Tensor(1)))
expectEqual(expected * 3, sumPullbackScalar(Tensor(3)))
// expectEqual(expected * 3, sumPullbackSqueezingAxes(Tensor(3)))
expectEqual(expected * 3, sumPullbackSqueezingAxes(Tensor(3)))
expectEqual(expected * 3, sumPullbackAlongAxes(Tensor(3)))
}

TensorADTests.testAllBackends("mean") {
let meanGradScalar = gradient { (a: Tensor<Float>) in a.mean() }
// let meanGradSqueezingAxes = gradient { (a: Tensor<Float>) in a.mean(squeezingAxes: 0, 1) }
let meanGradSqueezingAxes = gradient { (a: Tensor<Float>) in a.mean(squeezingAxes: 0, 1) }
let meanGradAlongAxes = gradient { (a: Tensor<Float>) in a.mean(alongAxes: 0, 1) }

let input = Tensor<Float>(ones: [2, 2])
let expected = Tensor<Float>(repeating: 0.25, shape: [2, 2])
expectEqual(expected, meanGradScalar(input))
// expectEqual(expected, meanGradSqueezingAxes(input))
expectEqual(expected, meanGradSqueezingAxes(input))
expectEqual(expected, meanGradAlongAxes(input))
}

TensorADTests.testAllBackends("variance") {
let varianceGradScalar = gradient { (a: Tensor<Float>) in a.variance() }
// let varianceGradSqueezingAxes = gradient { (a: Tensor<Float>) in a.variance(squeezingAxes: 0, 1) }
let varianceGradSqueezingAxes = gradient { (a: Tensor<Float>) in a.variance(squeezingAxes: 0, 1) }
let varianceGradAlongAxes = gradient { (a: Tensor<Float>) in a.variance(alongAxes: 0, 1) }

let input: Tensor<Float> = [[1, 2], [3, 4]]
let expected: Tensor<Float> = [[-0.75, -0.25], [0.25, 0.75]]
expectEqual(expected, varianceGradScalar(input))
// expectEqual(expected, varianceGradSqueezingAxes(input))
expectEqual(expected, varianceGradSqueezingAxes(input))
expectEqual(expected, varianceGradAlongAxes(input))
}

Expand Down