def METHOD_NAME():
    @task
    @annotate({"in_val": float, "return": {"squared": float}})
    def square(in_val):
        return in_val**2
    res = square(in_val=2.0)()
    assert res.output.squared == 4.0