             self.cell = cell
             self.hidden = hidden
         }
+    }
+
+    /// Returns a zero-valued state with shape compatible with the provided input.
+    public func zeroState(for input: Tensor<Scalar>) -> State {
+        let hiddenSize = fusedWeight.shape[1] / 4
+        return State(
+            cell: Tensor(zeros: [input.shape[0], hiddenSize]),
+            hidden: Tensor(zeros: [input.shape[0], hiddenSize]))
     }
 
     /// Returns the output obtained from applying the layer to the given input.

    associatedtype TimeStepOutput: Differentiable
    /// The state that may be preserved across time steps.
    associatedtype State: Differentiable
    /// The zero state.
    var zeroState: State { get }
}

public extension RNNCell {

        self.cell = cell()
    }

    @differentiable(wrt: (self, inputs), vjp: _vjpCallAsFunction(_:initialState:))
    public func callAsFunction(
        _ inputs: [Cell.TimeStepInput],
        initialState: Cell.State
    ) -> [Cell.TimeStepOutput] {
        if inputs.isEmpty { return [Cell.TimeStepOutput]() }
        var currentHiddenState = initialState
        var timeStepOutputs: [Cell.TimeStepOutput] = []
        for timeStepInput in inputs {
            let output = cell(input: timeStepInput, state: currentHiddenState)
            currentHiddenState = output.state
            timeStepOutputs.append(output.output)
        }
        return timeStepOutputs
    }

    @differentiable(wrt: (self, inputs))
    public func call(
        _ inputs: [Cell.TimeStepInput],
        initialState: Cell.State
    ) -> [Cell.TimeStepOutput] {
        callAsFunction(inputs, initialState: initialState)
    }

    @usableFromInline
