Skip to content
This repository was archived by the owner on Jul 1, 2023. It is now read-only.

Added a couple missing 'Tensor.moments' functions. #363

Merged
merged 6 commits into from
Jul 14, 2019
Merged

Added a couple missing 'Tensor.moments' functions. #363

Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
f437140
Added a couple missing 'Tensor.moments' functions.
eaplatanios Jul 14, 2019
7d45edc
Bug fix.
eaplatanios Jul 14, 2019
bb3f849
Bug fix.
eaplatanios Jul 14, 2019
b01b6a4
Bug fix.
eaplatanios Jul 14, 2019
507bf94
Bug fix.
eaplatanios Jul 14, 2019
d70b829
Added preconditions in the doc strings.
eaplatanios Jul 14, 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
54 changes: 48 additions & 6 deletions Sources/TensorFlow/Operators/Math.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2315,18 +2315,40 @@ public struct Moments<Scalar: TensorFlowFloatingPoint>: Differentiable {
public extension Tensor where Scalar: TensorFlowFloatingPoint {
/// Returns the mean and variance of this tensor along the specified axes. The reduced
/// dimensions are removed.
Copy link
Contributor

Choose a reason for hiding this comment

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

Add a precondition which states that axes.rank must be 0.

Copy link
Contributor Author

Choose a reason for hiding this comment

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

It has to be rank 1, but I didn't add it because it would cause lazy tensor materialization. The same is true for all reduction ops and we don't use any preconditions there.

Copy link
Contributor

Choose a reason for hiding this comment

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

Oh sorry I meant rank 1. Precondition should be added in comments at least.

Copy link
Contributor Author

Choose a reason for hiding this comment

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

Good point. I added it in the documentation strings.

///
/// - Parameter axes: The dimensions to reduce.
/// - Precondition: `axes` must have rank `1`.
/// - Precondition: Each value in `axes` must be in the range `-rank..<rank`.
@inlinable
@differentiable(wrt: self)
func moments(squeezingAxes axes: [Int]) -> Moments<Scalar> {
func moments(squeezingAxes axes: Tensor<Int32>) -> Moments<Scalar> {
let mean = self.mean(alongAxes: axes)
let variance = squaredDifference(self, mean).mean(alongAxes: axes)
return Moments<Scalar>(
mean: mean.squeezingShape(at: axes),
variance: variance.squeezingShape(at: axes))
let variance = squaredDifference(self, mean).mean(squeezingAxes: axes)
return Moments(
// The following is required because `Tensor.squeezingShape(at:)` does not accept
// `Tensor<Int32>`-valued arguments.
mean: mean.sum(squeezingAxes: axes),
variance: variance)
}

/// Returns the mean and variance of this tensor 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
@differentiable(wrt: self)
func moments(squeezingAxes axes: [Int]) -> Moments<Scalar> {
// TODO(TF-433): Remove workaround for differentiating `map`.
let axes = {axes.map(Int32.init)}()
return moments(squeezingAxes: Tensor<Int32>(axes))
}

/// Returns the mean and variance of this tensor 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
@differentiable(wrt: self)
func moments(squeezingAxes axes: Int...) -> Moments<Scalar> {
Expand All @@ -2342,16 +2364,36 @@ public extension Tensor where Scalar: TensorFlowFloatingPoint {

/// Returns the mean and variance of this tensor along the specified axes. The reduced
/// dimensions are retained with value `1`.
///
/// - Parameter axes: The dimensions to reduce.
/// - Precondition: `axes` must have rank `1`.
/// - Precondition: Each value in `axes` must be in the range `-rank..<rank`.
@inlinable
@differentiable(wrt: self)
func moments(alongAxes axes: [Int]) -> Moments<Scalar> {
func moments(alongAxes axes: Tensor<Int32>) -> Moments<Scalar> {
let mean = self.mean(alongAxes: axes)
let variance = squaredDifference(self, mean).mean(alongAxes: axes)
return Moments<Scalar>(mean: mean, variance: variance)
}

/// Returns the mean and variance of this tensor 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
@differentiable(wrt: self)
func moments(alongAxes axes: [Int]) -> Moments<Scalar> {
// TODO(TF-433): Remove workaround for differentiating `map`.
let axes = {axes.map(Int32.init)}()
return moments(alongAxes: Tensor<Int32>(axes))
}

/// Returns the mean and variance of this tensor 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
@differentiable(wrt: self)
func moments(alongAxes axes: Int...) -> Moments<Scalar> {
Expand Down