def dot_product(weights, features):
    return jnp.dot(weights, features)

        __
        __