Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / activation_memory / __init__.py: 100%

1# Copyright 2026 Huawei Technologies Co., Ltd 

2# 

3# Licensed under the Apache License, Version 2.0 (the "License"); 

4# you may not use this file except in compliance with the License. 

5# You may obtain a copy of the License at 

6# 

7# http://www.apache.org/licenses/LICENSE-2.0 

8# 

9# Unless required by applicable law or agreed to in writing, software 

10# distributed under the License is distributed on an "AS IS" BASIS, 

11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 

12# See the License for the specific language governing permissions and 

13# limitations under the License. 

14# ============================================================================ 

15"""Activation checkpointing related interfaces (PyTorch backend only).""" 

16from .checkpoint import CheckpointError 

17from .pinned_memory_pool import PinnedMemoryPool 

18from .recompute_state import get_recompute_state, is_recomputing 

19from .swap import SwapManager 

20from .wrapper import ( 

21 ActivationWrapper, 

22 AsyncSaveOnCpu, 

23 CheckpointExcludeWrapper, 

24 CheckpointWrapper, 

25 SwapWrapper, 

26 base_check_fn, 

27 ckpt_wrapper, 

28) 

29from .api import ( 

30 CheckpointPolicy, 

31 clear_recompute_session, 

32 checkpoint, 

33 checkpoint_exclude_wrapper, 

34 checkpoint_wrapper, 

35 recompute_handle, 

36 recompute_handle_collector_ctx, 

37 recompute_session_ctx, 

38 swap, 

39 swap_tensor_wrapper, 

40 swap_wrapper, 

41) 

42 

43__all__ = [ 

44 "ActivationWrapper", 

45 "AsyncSaveOnCpu", 

46 "CheckpointError", 

47 "CheckpointExcludeWrapper", 

48 "CheckpointPolicy", 

49 "CheckpointWrapper", 

50 "PinnedMemoryPool", 

51 "SwapManager", 

52 "SwapWrapper", 

53 "base_check_fn", 

54 "clear_recompute_session", 

55 "checkpoint", 

56 "checkpoint_exclude_wrapper", 

57 "checkpoint_wrapper", 

58 "ckpt_wrapper", 

59 "get_recompute_state", 

60 "is_recomputing", 

61 "recompute_handle", 

62 "recompute_handle_collector_ctx", 

63 "recompute_session_ctx", 

64 "swap", 

65 "swap_tensor_wrapper", 

66 "swap_wrapper", 

67]