         return TensorShape([1, weight.shape[1]])
     }
 
-    public var zeroState: Tensor<Scalar> {
-        return Tensor(zeros: stateShape)
-    }
-
-    public typealias State = Tensor<Scalar>
+    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 TimeStepInput = Tensor<Scalar>
     public typealias TimeStepOutput = State
     public typealias Input = RNNCellInput<TimeStepInput, State>

    /// - 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)
        return Output(output: newState, state: newState)
    }
}

                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.bias = Tensor(zeros: [hiddenSize])
    }

