def predict(params, input):
    """
    Description:

    This is the model to be trained. Pay close
    attention to the signature of this function. We use functional approach
    in JAX, unlike other libraries.

    input -- an scalar
    params -- a pytree data structure consists of parameters to be trained
    """
    w, b = params.w, params.b
    return w*input + b