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

Workaround for SR-10697. #123

Merged
merged 2 commits into from
May 16, 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
21 changes: 14 additions & 7 deletions Sources/DeepLearning/Layer.swift
Original file line number Diff line number Diff line change
Expand Up @@ -1299,11 +1299,19 @@ public struct SimpleRNNCell<Scalar: TensorFlowFloatingPoint>: RNNCell, VectorNum
return TensorShape([1, weight.shape[1]])
}

public var zeroState: Tensor<Scalar> {
return Tensor(zeros: stateShape)
public var zeroState: State {
return State(Tensor(zeros: stateShape))
}

// TODO(TF-507): Revert to `typealias State = Tensor<Scalar>` after
// SR-10697 is fixed.
public struct State: Differentiable {
public let value: Tensor<Scalar>
public init(_ value: Tensor<Scalar>) {
self.value = value
}
}

public typealias State = Tensor<Scalar>
public typealias TimeStepInput = Tensor<Scalar>
public typealias TimeStepOutput = State
public typealias Input = RNNCellInput<TimeStepInput, State>
Expand All @@ -1319,8 +1327,7 @@ public struct SimpleRNNCell<Scalar: TensorFlowFloatingPoint>: RNNCell, VectorNum
seed: (Int64, Int64) = (Int64.random(in: Int64.min..<Int64.max),
Int64.random(in: Int64.min..<Int64.max))) {
let concatenatedInputSize = inputSize + hiddenSize
self.weight = Tensor(glorotUniform: [concatenatedInputSize, hiddenSize],
seed: seed)
self.weight = Tensor(glorotUniform: [concatenatedInputSize, hiddenSize], seed: seed)
self.bias = Tensor(zeros: [hiddenSize])
}

Expand All @@ -1330,8 +1337,8 @@ public struct SimpleRNNCell<Scalar: TensorFlowFloatingPoint>: RNNCell, VectorNum
/// - Returns: The hidden state.
@differentiable
public func call(_ input: Input) -> Output {
let concatenatedInput = input.input.concatenated(with: input.state, alongAxis: 1)
let newState = tanh(matmul(concatenatedInput, weight) + bias)
let concatenatedInput = input.input.concatenated(with: input.state.value, alongAxis: 1)
let newState = State(tanh(matmul(concatenatedInput, weight) + bias))
return Output(output: newState, state: newState)
}
}
Expand Down