         let learningRate = self.learningRate * 1 / (1 + decay * step)
         // Note: `stepSize` is split into two lines to avoid the "compiler is unable to type-check
         // this expression in reasonable time" error.
-        var stepSize = learningRate * sqrt(1 - pow(beta2, step))
-        stepSize = stepSize / (1 - pow(beta1, step))
+        var stepSize = learningRate * sqrtf(1 - powf(beta2, step))
+        stepSize = stepSize / (1 - powf(beta1, step))
         firstMoments = firstMoments * beta1 + direction * (1 - beta1)
 
         // Update `infinityNorm` using a key path approach because `max(_:_:)` cannot be 

        let learningRate = self.learningRate * 1 / (1 + decay * step)
        // Note: `stepSize` and `secondMoments` are split into two lines to avoid the "compiler is 
        // unable to type-check this expression in reasonable time" error.
        var stepSize = learningRate * sqrt(1 - pow(beta2, step))
        stepSize = stepSize / (1 - pow(beta1, step))
        firstMoments = firstMoments * beta1 + direction * (1 - beta1)
        secondMoments = secondMoments * beta2
        secondMoments += direction .* direction * (1 - beta2)

        if N_sma_t > 5 {
            // Compute bias-corrected second moments, rectification and adapted momentum.
            let secondMoments_h = Model.TangentVector.sqrt(secondMoments) + epsilon
            let stepSize = sqrtf(
                (N_sma_t - 4) * (N_sma_t - 2) * N_sma_inf / (
                     (N_sma_inf - 4) * (N_sma_inf - 2) * (N_sma_t)
                ))
            model.move(along: -stepSize * sqrtf(1 - beta2Power) * firstMoments ./ secondMoments_h)
        } else {
            // Update with un-adapted momentum.
            let stepSize = self.learningRate * step / (1 - beta1Power)
