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

Minor tweak to allow min and max to be differentiable. #159

Merged
merged 1 commit into from
May 31, 2019
Merged
Changes from all commits
Commits
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
12 changes: 10 additions & 2 deletions Sources/TensorFlow/Operators/Math.swift
Original file line number Diff line number Diff line change
Expand Up @@ -915,6 +915,7 @@ public extension Tensor where Scalar: Numeric & Comparable {
// NOTE: This overload is necessary, otherwise `min()` would refer to the variadic method
// `min(squeezingAxes:)` with zero indices.
@inlinable
@differentiable(where Scalar: TensorFlowFloatingPoint)
func min() -> Tensor {
let axes = Tensor<Int32>(rangeFrom: 0, to: Int32(rank), stride: 1)
return min(squeezingAxes: axes)
Expand All @@ -923,6 +924,7 @@ public extension Tensor where Scalar: Numeric & Comparable {
// NOTE: This overload is necessary, otherwise `max()` would refer to the variadic method
// `max(squeezingAxes:)` with zero indices.
@inlinable
@differentiable(where Scalar: TensorFlowFloatingPoint)
func max() -> Tensor {
let axes = Tensor<Int32>(rangeFrom: 0, to: Int32(rank), stride: 1)
return max(squeezingAxes: axes)
Expand All @@ -943,15 +945,18 @@ public extension Tensor where Scalar: Numeric & Comparable {
/// - Parameter axes: The dimensions to reduce.
/// - Precondition: Each value in `axes` must be in the range `-rank..<rank`.
@inlinable
@differentiable(wrt: self where Scalar: TensorFlowFloatingPoint)
func max(squeezingAxes axes: [Int]) -> Tensor {
let axes = axes.map(Int32.init)
// TODO(TF-433): Remove workaround for differentiating `map`.
let axes = {axes.map(Int32.init)}()
return max(squeezingAxes: Tensor<Int32>(axes))
}

/// Returns the maximum values 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 where Scalar: TensorFlowFloatingPoint)
func max(squeezingAxes axes: Int...) -> Tensor {
return max(squeezingAxes: axes)
}
Expand All @@ -971,15 +976,18 @@ public extension Tensor where Scalar: Numeric & Comparable {
/// - Parameter axes: The dimensions to reduce.
/// - Precondition: Each value in `axes` must be in the range `-rank..<rank`.
@inlinable
@differentiable(wrt: self where Scalar: TensorFlowFloatingPoint)
func min(squeezingAxes axes: [Int]) -> Tensor {
let axes = axes.map(Int32.init)
// TODO(TF-433): Remove workaround for differentiating `map`.
let axes = {axes.map(Int32.init)}()
return min(squeezingAxes: Tensor<Int32>(axes))
}

/// Returns the minimum values 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 where Scalar: TensorFlowFloatingPoint)
func min(squeezingAxes axes: Int...) -> Tensor {
return min(squeezingAxes: axes)
}
Expand Down