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

Makes 'Tensor.gathering' generic over 'Int32' and 'Int64'. #272

Merged
merged 3 commits into from
Jun 21, 2019
Merged

Makes 'Tensor.gathering' generic over 'Int32' and 'Int64'. #272

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
68828a6
Makes 'Tensor.gathering' generic over 'Int32' and 'Int64'. Fixes #259.
eaplatanios Jun 21, 2019
798afb8
Added a 'TensorFlowIndex' protocol.
eaplatanios Jun 21, 2019
28f21ea
Merge remote-tracking branch 'upstream/master' into gathering-int32-i…
eaplatanios Jun 21, 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
7 changes: 7 additions & 0 deletions Sources/TensorFlow/Core/DataTypes.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -70,6 +70,13 @@ public typealias TensorFlowNumeric = TensorFlowScalar & Numeric
public typealias TensorFlowSignedNumeric = TensorFlowScalar & SignedNumeric
public typealias TensorFlowInteger = TensorFlowScalar & BinaryInteger

/// An integer data type that represents integer types which can be used as tensor indices in
/// TensorFlow.
public protocol TensorFlowIndex: TensorFlowInteger {}

extension Int32: TensorFlowIndex {}
extension Int64: TensorFlowIndex {}

/// A floating-point data type that conforms to `Differentiable` and is compatible with TensorFlow.
///
/// - Note: `Tensor` conditionally conforms to `Differentiable` when the `Scalar` associated type
Expand Down
23 changes: 13 additions & 10 deletions Sources/TensorFlow/Operators/Basic.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -383,7 +383,10 @@ public extension Tensor {
/// - Returns: The gathered tensor.
@inlinable
@differentiable(wrt: self, vjp: _vjpGathering where Scalar : TensorFlowFloatingPoint)
func gathering(atIndices indices: Tensor<Int32>, alongAxis axis: Int = 0) -> Tensor {
func gathering<Index: TensorFlowIndex>(
atIndices indices: Tensor<Index>,
alongAxis axis: Int = 0
) -> Tensor {
return Raw.gatherV2(params: self, indices: indices, axis: Tensor<Int32>(Int32(axis)))
}

Expand All @@ -408,15 +411,15 @@ public extension Tensor {
/// - Returns: The gathered tensor.
@inlinable
@differentiable(wrt: self where Scalar: TensorFlowFloatingPoint)
func batchGathering(atIndices indices: Tensor<Int32>) -> Tensor {
func batchGathering<Index: TensorFlowIndex>(atIndices indices: Tensor<Index>) -> Tensor {
var batchIndices = indices
var accumulated = Tensor<Int32>(ones: [])
accumulated *= Swift.withoutDerivative(at: shapeTensor) { $0[1] }
var accumulated = Tensor<Index>(ones: [])
accumulated *= Swift.withoutDerivative(at: shapeTensor) { Tensor<Index>($0[1]) }
let dValue = Swift.withoutDerivative(at: shapeTensor) { $0[0] }
let dIndices = Tensor<Int32>(
rangeFrom: Tensor<Int32>(zeros: []),
to: dValue,
stride: Tensor<Int32>(ones: [])
let dIndices = Tensor<Index>(
rangeFrom: Tensor<Index>(zeros: []),
to: Tensor<Index>(dValue),
stride: Tensor<Index>(ones: [])
) * accumulated
let dShape = Tensor<Int32>(concatenating: [
dValue.rankLifted(),
Expand Down Expand Up @@ -519,8 +522,8 @@ internal extension Tensor where Scalar: TensorFlowFloatingPoint {
}

@inlinable
func _vjpGathering(
atIndices indices: Tensor<Int32>,
func _vjpGathering<Index: TensorFlowIndex>(
atIndices indices: Tensor<Index>,
alongAxis axis: Int = 0
) -> (Tensor, (Tensor) -> Tensor) {
let result = gathering(atIndices: indices, alongAxis: axis)
Expand Down