     /// A small scalar added to the denominator to improve numerical stability.
     public var epsilon: Float
     /// The alpha values for all model differentiable variables.
-    public var alpha: Model.AllDifferentiableVariables
+    public var alpha: Model.TangentVector = .zero
 
     public init(
         for model: __shared Model,


/// AdaGrad optimizer.
///
/// Individually adapts the learning rates of all model parameters by scaling them inversely proportional to
/// the square root of the sum of all the historical squared values of the gradient.
///
/// Reference: ["Adaptive Subgradient Methods for Online Learning and Stochastic Optimization"](
/// http://www.jmlr.org/papers/volume12/duchi11a/duchi11a.pdf)
public class AdaGrad<Model: Layer>: Optimizer
    where Model.AllDifferentiableVariables == Model.TangentVector {
    public typealias Model = Model
    /// The learning rate.
    public var learningRate: Float

        self.learningRate = learningRate
        self.rho = rho
        self.epsilon = epsilon
    }

    public func update(_ model: inout Model, along direction: Model.TangentVector) {
        update(&model.allDifferentiableVariables, along: direction)
    }

    // TODO: Deprecate this when `Differentiable.AllDifferentiableVariables` is removed.
    public func update(
        _ model: inout Model.AllDifferentiableVariables,
        along direction: Model.TangentVector
    ) {
        alpha = rho + direction .* direction
        let denominator = Model.TangentVector.sqrt(alpha) + epsilon
        model.move(along: -learningRate * direction ./ denominator)
    }
}

