Function convert_to_enum

axlearn/common/utils.py:311–314

(p: Union[RematType, bool])

    combine_fn = maybe_instantiate(combine_fn)

    def convert_to_enum(p: Union[RematType, bool]) -> RematType:
        if isinstance(p, bool):
            p = Saveable if p else Recompute
        return p

    def policy(prim, *args, **kwargs):
        p1 = convert_to_enum(policy_1(prim, *args, **kwargs))