def reducer(arg1: TfVal, arg2: TfVal) -> TfVal:
    closed_jaxpr = core.ClosedJaxpr(jaxpr, consts)
    (res,) = _interpret_jaxpr(closed_jaxpr, arg1, arg2)
    return res
def reducer(arg1: TfVal, arg2: TfVal) -> TfVal:
    closed_jaxpr = core.ClosedJaxpr(jaxpr, consts)
    (res,) = _interpret_jaxpr(closed_jaxpr, arg1, arg2)
    return res