Skip to content

Add variance() and variance(squeezingAxes:). #23811

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 9, 2019
Merged

Add variance() and variance(squeezingAxes:). #23811

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
52f5910
Add `variance()` and `variance(squeezingAxes:)`.
dan-zheng Apr 5, 2019
40d73ab
Add derivatives for `squeezingAxes:` versions of `sum` and `mean`.
dan-zheng Apr 5, 2019
0c28110
Test variance.
dan-zheng Apr 5, 2019
8ce92df
Test variance derivatives.
dan-zheng Apr 5, 2019
d8aeafc
Fix `variance` derivatives.
dan-zheng Apr 9, 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
34 changes: 29 additions & 5 deletions stdlib/public/TensorFlow/Gradients.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -568,6 +568,11 @@ extension Tensor where Scalar : TensorFlowFloatingPoint {
//===----------------------------------------------------------------------===//

extension Tensor where Scalar : TensorFlowFloatingPoint {
@inlinable
func _vjpSum() -> (Tensor, (Tensor) -> Tensor) {
return (sum(), { [shape = shapeTensor] in $0.broadcast(toShape: shape) })
}

@inlinable
func _vjpMean() -> (Tensor, (Tensor) -> Tensor) {
return (mean(), { [shape = shapeTensor, count = scalarCountTensor] in
Expand All @@ -576,8 +581,15 @@ extension Tensor where Scalar : TensorFlowFloatingPoint {
}

@inlinable
func _vjpSum() -> (Tensor, (Tensor) -> Tensor) {
return (sum(), { [shape = shapeTensor] in $0.broadcast(toShape: shape) })
func _vjpSum(alongAxes axes: [Int32]) -> (Tensor, (Tensor) -> Tensor) {
let value = sum(alongAxes: axes)
return (value, { [shape = shapeTensor] in $0.broadcast(toShape: shape) })
}

@inlinable
func _vjpSum(squeezingAxes axes: [Int32]) -> (Tensor, (Tensor) -> Tensor) {
let value = sum(squeezingAxes: axes)
return (value, { [shape = shapeTensor] in $0.broadcast(toShape: shape) })
}

@inlinable
Expand All @@ -590,9 +602,21 @@ extension Tensor where Scalar : TensorFlowFloatingPoint {
}

@inlinable
func _vjpSum(alongAxes axes: [Int32]) -> (Tensor, (Tensor) -> Tensor) {
let value = sum(alongAxes: axes)
return (value, { [shape = shapeTensor] in $0.broadcast(toShape: shape) })
func _vjpMean(squeezingAxes axes: [Int32]) -> (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(alongAxes axes: [Int32]) -> (Tensor, (Tensor) -> Tensor) {
let value = mean(alongAxes: axes)
return (value, { [shape = shapeTensor,
count = axes.map { shape[$0] }.reduce(1, *)] in
$0.broadcast(toShape: shape) / Tensor(Scalar(count))
})
}
}

Expand Down
153 changes: 98 additions & 55 deletions stdlib/public/TensorFlow/Ops.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1251,18 +1251,6 @@ public extension Tensor where Scalar : Numeric & Comparable {
}

public extension Tensor where Scalar : Numeric {
// NOTE: This overload is necessary, otherwise `mean()` would refer
// to the variadic method `mean(squeezingAxes:)` with zero indices.
@differentiable(
wrt: self, vjp: _vjpMean()
where Scalar : TensorFlowFloatingPoint
)
@inlinable @inline(__always)
func mean() -> Tensor {
let axes = Tensor<Int32>(rangeFrom: 0, to: rank, stride: 1)
return Raw.mean(self, reductionIndices: axes)
}

// NOTE: This overload is necessary, otherwise `sum()` would refer
// to the variadic method `sum(squeezingAxes:)` with zero indices.
@inlinable @inline(__always)
Expand All @@ -1283,30 +1271,37 @@ public extension Tensor where Scalar : Numeric {
return Raw.prod(self, reductionIndices: axes)
}

/// Returns the arithmetic mean along the specified axes. The reduced
/// dimensions are removed.
/// - Parameter axes: The dimensions to reduce.
/// - Precondition: Each value in `axes` must be in the range `-rank...rank`.
// NOTE: This overload is necessary, otherwise `mean()` would refer
// to the variadic method `mean(squeezingAxes:)` with zero indices.
@differentiable(
wrt: self, vjp: _vjpMean()
where Scalar : TensorFlowFloatingPoint
)
@inlinable @inline(__always)
func mean(squeezingAxes axes: [Int32]) -> Tensor {
return Raw.mean(self, reductionIndices: Tensor<Int32>(axes),
keepDims: false)
func mean() -> Tensor {
let axes = Tensor<Int32>(rangeFrom: 0, to: rank, stride: 1)
return Raw.mean(self, reductionIndices: axes)
}

/// Returns the arithmetic mean along the specified axes. The reduced
/// dimensions are removed.
/// - Parameter axes: The dimensions to reduce.
/// - Precondition: Each value in `axes` must be in the range `-rank...rank`.
// NOTE: This overload is necessary, otherwise `mean()` would refer
// to the variadic method `mean(squeezingAxes:)` with zero indices.
@differentiable(wrt: self where Scalar : TensorFlowFloatingPoint)
@inlinable @inline(__always)
func mean(squeezingAxes axes: Int32...) -> Tensor {
return mean(squeezingAxes: axes)
func variance() -> Tensor {
let mean = self.mean()
let squaredDiff = (self - mean).squared()
return squaredDiff.mean()
}

/// Returns the sum along the specified axes. The reduced dimensions are
/// removed.
/// - Parameter axes: The dimensions to reduce.
/// - Precondition: Each value in `axes` must be in the range `-rank...rank`.
@inlinable @inline(__always)
@differentiable(
wrt: self, vjp: _vjpSum(squeezingAxes:)
where Scalar : TensorFlowFloatingPoint
)
func sum(squeezingAxes axes: [Int32]) -> Tensor {
return Raw.sum(self, reductionIndices: Tensor<Int32>(axes), keepDims: false)
}
Expand Down Expand Up @@ -1340,36 +1335,48 @@ public extension Tensor where Scalar : Numeric {
}

/// Returns the arithmetic mean along the specified axes. The reduced
/// dimensions are retained with value 1.
/// dimensions are removed.
/// - Parameter axes: The dimensions to reduce.
/// - Precondition: Each value in `axes` must be in the range `-rank..<rank`.
/// - Precondition: Each value in `axes` must be in the range `-rank...rank`.
@inlinable @inline(__always)
@differentiable(
wrt: self, vjp: _vjpMean(alongAxes:)
wrt: self, vjp: _vjpMean(squeezingAxes:)
where Scalar : TensorFlowFloatingPoint
)
func mean(alongAxes axes: Tensor<Int32>) -> Tensor {
return Raw.mean(self, reductionIndices: axes, keepDims: true)
func mean(squeezingAxes axes: [Int32]) -> Tensor {
return Raw.mean(self, reductionIndices: Tensor<Int32>(axes),
keepDims: false)
}

/// Returns the arithmetic mean along the specified axes. The reduced
/// dimensions are retained with value 1.
/// dimensions are removed.
/// - Parameter axes: The dimensions to reduce.
/// - Precondition: Each value in `axes` must be in the range `-rank...rank`.
@inlinable @inline(__always)
func mean(squeezingAxes axes: Int32...) -> Tensor {
return mean(squeezingAxes: axes)
}

/// Returns the variance along the specified axes. The reduced dimensions are
/// retained with value 1. Does not apply Bessel's correction.
/// - Parameter axes: The dimensions to reduce.
/// - Precondition: Each value in `axes` must be in the range `-rank..<rank`.
@inlinable @inline(__always)
@differentiable(wrt: self where Scalar : TensorFlowFloatingPoint)
func mean(alongAxes axes: [Int32]) -> Tensor {
return mean(alongAxes: Tensor<Int32>(axes))
func variance(squeezingAxes axes: Int32...) -> Tensor {
return variance(squeezingAxes: axes)
}

/// Returns the arithmetic mean along the specified axes. The reduced
/// dimensions are retained with value 1.
/// Returns the variance along the specified axes. The reduced dimensions are
/// removed. Does not apply Bessel's correction.
/// - Parameter axes: The dimensions to reduce.
/// - Precondition: Each value in `axes` must be in the range `-rank..<rank`.
@inlinable @inline(__always)
@differentiable(wrt: self where Scalar : TensorFlowFloatingPoint)
func mean(alongAxes axes: Int32...) -> Tensor {
return mean(alongAxes: axes)
func variance(squeezingAxes axes: [Int32]) -> Tensor {
let mean = self.mean(alongAxes: axes)
let squaredDiff = (self - mean).squared()
return squaredDiff.mean(squeezingAxes: axes)
}

/// Returns the sum along the specified axes. The reduced dimensions are
Expand All @@ -1395,6 +1402,60 @@ public extension Tensor where Scalar : Numeric {
return sum(alongAxes: axes)
}

/// Returns the product along the specified axes. The reduced dimensions are
/// retained with value 1.
/// - Parameter axes: The dimensions to reduce.
/// - Precondition: Each value in `axes` must be in the range `-rank..<rank`.
@inlinable @inline(__always)
func product(alongAxes axes: [Int32]) -> Tensor {
return Raw.prod(self, reductionIndices: Tensor<Int32>(axes), keepDims: true)
}

/// Returns the product along the specified axes. The reduced dimensions are
/// retained with value 1.
/// - Parameter axes: The dimensions to reduce.
/// - Precondition: Each value in `axes` must be in the range `-rank..<rank`.
@inlinable @inline(__always)
func product(alongAxes axes: Int32...) -> Tensor {
return product(alongAxes: axes)
}

/// Returns the arithmetic mean along the specified axes. The reduced
/// dimensions are retained with value 1.
/// - Parameter axes: The dimensions to reduce.
/// - Precondition: Each value in `axes` must be in the range `-rank..<rank`.
@inlinable @inline(__always)
@differentiable(
wrt: self, vjp: _vjpMean(alongAxes:)
where Scalar : TensorFlowFloatingPoint
)
func mean(alongAxes axes: Tensor<Int32>) -> Tensor {
return Raw.mean(self, reductionIndices: axes, keepDims: true)
}

/// Returns the arithmetic mean along the specified axes. The reduced
/// dimensions are retained with value 1.
/// - Parameter axes: The dimensions to reduce.
/// - Precondition: Each value in `axes` must be in the range `-rank..<rank`.
@inlinable @inline(__always)
@differentiable(
wrt: self, vjp: _vjpMean(alongAxes:)
where Scalar : TensorFlowFloatingPoint
)
func mean(alongAxes axes: [Int32]) -> Tensor {
return mean(alongAxes: Tensor<Int32>(axes))
}

/// Returns the arithmetic mean along the specified axes. The reduced
/// dimensions are retained with value 1.
/// - Parameter axes: The dimensions to reduce.
/// - Precondition: Each value in `axes` must be in the range `-rank..<rank`.
@inlinable @inline(__always)
@differentiable(wrt: self where Scalar : TensorFlowFloatingPoint)
func mean(alongAxes axes: Int32...) -> Tensor {
return mean(alongAxes: axes)
}

/// Returns the variance along the specified axes. The reduced dimensions are
/// retained with value 1. Does not apply Bessel's correction.
/// - Parameter axes: The dimensions to reduce.
Expand Down Expand Up @@ -1426,24 +1487,6 @@ public extension Tensor where Scalar : Numeric {
func variance(alongAxes axes: [Int32]) -> Tensor {
return variance(alongAxes: Tensor<Int32>(axes))
}

/// Returns the product along the specified axes. The reduced dimensions are
/// retained with value 1.
/// - Parameter axes: The dimensions to reduce.
/// - Precondition: Each value in `axes` must be in the range `-rank..<rank`.
@inlinable @inline(__always)
func product(alongAxes axes: [Int32]) -> Tensor {
return Raw.prod(self, reductionIndices: Tensor<Int32>(axes), keepDims: true)
}

/// Returns the product along the specified axes. The reduced dimensions are
/// retained with value 1.
/// - Parameter axes: The dimensions to reduce.
/// - Precondition: Each value in `axes` must be in the range `-rank..<rank`.
@inlinable @inline(__always)
func product(alongAxes axes: Int32...) -> Tensor {
return product(alongAxes: axes)
}
}

//===----------------------------------------------------------------------===//
Expand Down
39 changes: 31 additions & 8 deletions test/TensorFlowRuntime/tensor.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -178,14 +178,37 @@ TensorTests.testAllBackends("Reduction") {
#if !TPU
// 2 x 5
let x = Tensor<Float>([[1, 2, 3, 4, 5], [1, 2, 3, 4, 5]])
expectEqual(ShapedArray(shape: [5], scalars: [2, 4, 6, 8, 10]),
x.sum(squeezingAxes: 0).toHost(shape: []).array)
expectEqual(ShapedArray(shape: [1, 5], scalars: [2, 4, 6, 8, 10]),
x.sum(alongAxes: 0).toHost(shape: []).array)
expectEqual(ShapedArray(shape: [5], scalars: [1, 4, 9, 16, 25]),
x.product(squeezingAxes: 0).toHost(shape: []).array)
expectEqual(ShapedArray(shape: [1, 5], scalars: [1, 4, 9, 16, 25]),
x.product(alongAxes: 0).toHost(shape: []).array)
expectEqual(Tensor(30), x.sum().toHost(shape: []))
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I wonder if reduction ops still fail on TPU. The #if !TPU condition was added months ago (b/111815968).

cc @pschuh @bgogul: I wonder if you have any clue if this test still fails on TPU?
Or if toHost(shape: ...) is necessary everywhere?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Also, the arguments to toHost(shape:) don't seem quite right (it shouldn't always be []).

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Update: I'll remove the toHost calls to force us to re-evaluate this bug.

expectEqual(Tensor(shape: [5], scalars: [2, 4, 6, 8, 10]),
x.sum(squeezingAxes: 0).toHost(shape: []))
expectEqual(Tensor(shape: [1, 5], scalars: [2, 4, 6, 8, 10]),
x.sum(alongAxes: 0).toHost(shape: []))

expectEqual(Tensor(14400), x.product().toHost(shape: []))
expectEqual(Tensor(shape: [5], scalars: [1, 4, 9, 16, 25]),
x.product(squeezingAxes: 0).toHost(shape: []))
expectEqual(Tensor(shape: [1, 5], scalars: [1, 4, 9, 16, 25]),
x.product(alongAxes: 0).toHost(shape: []))

expectEqual(Tensor(3), x.mean().toHost(shape: []))
expectEqual(Tensor(shape: [5], scalars: [1, 2, 3, 4, 5]),
x.mean(squeezingAxes: 0).toHost(shape: []))
expectEqual(Tensor(shape: [5], scalars: [1, 2, 3, 4, 5]),
x.mean(alongAxes: 0).toHost(shape: []))
expectEqual(Tensor(shape: [2], scalars: [3, 3]),
x.mean(squeezingAxes: 1).toHost(shape: []))
expectEqual(Tensor(shape: [1, 2], scalars: [3, 3]),
x.mean(alongAxes: 1).toHost(shape: []))

expectEqual(Tensor(2), x.variance().toHost(shape: []))
expectEqual(Tensor(shape: [5], scalars: [0, 0, 0, 0, 0]),
x.variance(squeezingAxes: 0).toHost(shape: []))
expectEqual(Tensor(shape: [5], scalars: [0, 0, 0, 0, 0]),
x.variance(alongAxes: 0).toHost(shape: []))
expectEqual(Tensor(shape: [2], scalars: [2, 2]),
x.variance(squeezingAxes: 1).toHost(shape: []))
expectEqual(Tensor(shape: [1, 2], scalars: [2, 2]),
x.variance(alongAxes: 1).toHost(shape: []))
#endif // !TPU
}

Expand Down
12 changes: 12 additions & 0 deletions test/TensorFlowRuntime/tensor_autodiff_runtime.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -121,6 +121,18 @@ TensorADTests.testAllBackends("mean") {
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 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, varianceGradAlongAxes(input))
}

TensorADTests.testAllBackends("expandingShape") {
let f1 = { (a: Tensor<Float>) in a.expandingShape(at: 0).squared() }
let f2 = { (a: Tensor<Float>) in a.squared().expandingShape(at: 0) }
Expand Down