def embedding_checkpoint_wrapper(
    activation_checkpointing: Union[str, Dict],
) -> Callable:
    """return encoder embedding activation checkpoint wrapper"""
    validate_checkpointing_config(activation_checkpointing)

    if isinstance(activation_checkpointing, str):
        if activation_checkpointing:
            if activation_checkpointing == "offload":
                return offload_wrapper
            return partial(checkpoint_wrapper)
        return lambda x: x

    if isinstance(activation_checkpointing, dict):
        enabled = activation_checkpointing.get("embed", False)
        if enabled:
            offloading = activation_checkpointing.get("offload", False)
            if offloading:
                return offload_wrapper
            impl = (
                CheckpointImpl.REENTRANT
                if activation_checkpointing.get("reentrant", False)
                else CheckpointImpl.NO_REENTRANT
            )
            return partial(checkpoint_wrapper, checkpoint_impl=impl)
        return lambda x: x
    raise ValueError("Invalid activation_checkpointing config")
return encoder embedding activation checkpoint wrapper
tests/cpu/hf_configs/phi4/speech_conformer_encoder.py