Skip to content

Commit

Permalink
[stdlib] Improve _horner_evaluate docstring
Browse files Browse the repository at this point in the history
- Small simplification to the function `_horner_evaluate`.
- Improve its docstring to match the argument names.
- Finish off removing `unroll` by cleaning up the imports.

Signed-off-by: Yiwu Chen <[email protected]>
  • Loading branch information
soraros committed Jun 20, 2024
1 parent 279ade2 commit 15322f7
Showing 1 changed file with 8 additions and 13 deletions.
21 changes: 8 additions & 13 deletions stdlib/src/math/polynomial.mojo
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,6 @@ from math.polynomial import polynomial_evaluate

from collections import List

from utils.loop import unroll

# ===----------------------------------------------------------------------===#
# polynomial_evaluate
Expand Down Expand Up @@ -64,11 +63,12 @@ fn _horner_evaluate[
](x: SIMD[dtype, simd_width]) -> SIMD[dtype, simd_width]:
"""Evaluates the polynomial using the passed in value and the specified
coefficients using the Horner scheme. The Horner scheme evaluates the
polynomial as `horner(val, coeffs)` where val is a scalar and coeffs is a
list of coefficients [c0, c1, c2, ..., cn] by:
polynomial at point x as `horner(x, coeffs)` where x is a scalar and coeffs
is a list of coefficients [c0, c1, c2, ..., cn] by:
```
horner(val, coeffs) = c0 + val * (c1 + val * (c2 + val * (... + val * cn)))
= fma(val, horner(val, coeffs[1:]), c0)
horner(x, coeffs)
= c0 + x * (c1 + x * (c2 + x * (... + x * cn)))
= fma(x, horner(x, coeffs[1:]), coeffs[0])
```
Parameters:
Expand All @@ -80,17 +80,12 @@ fn _horner_evaluate[
x: The value to compute the polynomial with.
Returns:
The polynomial evaluation results using the specified value and the
constant coefficients.
The polynomial specified by the coefficients evaluated at value x.
"""
alias num_coefficients = len(coefficients)
alias c_last = coefficients[num_coefficients - 1]
alias c_second_from_last = coefficients[num_coefficients - 2]

var result = x.fma(c_last, c_second_from_last)
var result = x.fma(coefficients[-1], coefficients[-2])

@parameter
for i in reversed(range(num_coefficients - 2)):
for i in reversed(range(len(coefficients) - 2)):
alias c = coefficients[i]
result = result.fma(x, c)

Expand Down

0 comments on commit 15322f7

Please sign in to comment.