Compute polynomials efficiently with numpy and pytorch (differentiable).

polynomial.py

import torch

def polynomial(coefficients, x):
    """Evaluate polynomials using Horner method.

    The coefficients are from highest to lowest order.

    Args:
        coefficients: Tensor of size (N, *K).
            K is any broadcastable size to `x.size()`.
        x: Arbitrary tensor.

    Returns:
        The evaluated polynomial at `x` with the same size.

    """
    out = x.new_zeros(x.size()) if torch.is_tensor(x) else 0
    coefficients = iter(coefficients)
    out += next(coefficients, 0)
    for c in coefficients:
        out *= x
        out += c
    return out