Coverage for tests/torch_mock.py: 95%

1""" 

2A mock for the torch module with specific mocked attributes and methods for testing purposes. 

3""" 

4 

5from unittest.mock import MagicMock 

6 

7import importlib.util 

8 

9from typing import Callable 

10from typing import Tuple 

11from typing import Dict 

12from typing import Any 

13from typing import List 

14from typing import Optional 

15 

16 

17class FakeModule: 

18 def __init__( 

19 self, 

20 *args: Any, 

21 **kwargs: Any 

22 ) -> None: 

23 pass 

24 

25 

26class FakeLayerNorm(FakeModule): 

27 def __init__( 

28 self, 

29 *args: Any, 

30 **kwargs: Any 

31 ) -> None: 

32 super().__init__(*args, **kwargs) 

33 

34 

35""" 

36class FakeFunctional(MagicMock): 

37 pass 

38 

39 

40class FakeNN: 

41 Module = FakeModule 

42 

43 def __getattr__(self, name): 

44 cls = type(name, (FakeModule,), {}) 

45 setattr(self, name, cls) 

46 return cls 

47""" 

48 

49 

50class TorchMock(MagicMock): 

51 """A mock for the torch module with specific mocked attributes and methods.""" 

52 

53 def __init__( 

54 self, 

55 *args: Tuple, 

56 **kwargs: Dict 

57 ) -> None: 

58 super().__init__(*args, **kwargs) 

59 

60 self.__spec__ = importlib.util.spec_from_loader("torch", loader=None) 

61 

62 # self.nn = FakeNN() 

63 

64 # Define real exception class for torch.OutOfMemoryError 

65 self.OutOfMemoryError = type("OutOfMemoryError", (RuntimeError,), {}) 

66 

67 # Define real type for torch.Tensor so isinstance() checks work 

68 self.Tensor = type("Tensor", (), {}) 

69 

70 # Mock torch.cuda.memory_allocated 

71 self.cuda = MagicMock() 

72 self.cuda.memory_allocated.return_value = 0 

73 self.cuda.device_count.return_value = 1 

74 self.cuda.get_device_name.side_effect = lambda device=None: f"MockDevice:{device if device is not None else 0}" 

75 self.cuda.synchronize = lambda device=None: None 

76 

77 # Mock @torch.inference_mode() decorator 

78 self.inference_mode = lambda: self._noop_decorator 

79 

80 # Mock tensor creation functions 

81 self.randn = lambda *shape, **kwargs: self._mock_tensor(*shape) 

82 self.ones = lambda *shape, **kwargs: self._mock_tensor(*shape, fill_value=1) 

83 self.tensor = lambda data, **kwargs: self._mock_tensor_from_data(data) 

84 

85 # Mock load/save 

86 self._save_store: Dict[Any, Any] = {} # internal dict to track saved "objects" 

87 

88 def mock_save( 

89 obj: Any, 

90 f: Any, 

91 **kwargs: Dict[str, Any] 

92 ) -> None: 

93 self._save_store[f] = obj 

94 return None 

95 

96 def mock_load( 

97 f: Any, 

98 map_location: Optional[Any] = None, 

99 **kwargs: Dict[str, Any] 

100 ) -> Any: 

101 return self._save_store.get(f, MagicMock(name="LoadedTensor")) 

102 

103 self.save = mock_save 

104 self.load = mock_load 

105 

106 @staticmethod 

107 def _noop_decorator(func: Callable) -> Callable: 

108 return func 

109 

110 @staticmethod 

111 def _mock_tensor(*shape: int, fill_value: Any = 0) -> MagicMock: 

112 m = MagicMock() 

113 m.shape = shape 

114 m.ndim = len(shape) 

115 m.fill_value = fill_value 

116 

117 def chunk(chunks: int, dim: int = 0) -> List[Any]: 

118 return [m for _ in range(chunks)] 

119 

120 m.chunk.side_effect = chunk 

121 m.__getitem__.side_effect = lambda idx: fill_value 

122 return m 

123 

124 @staticmethod 

125 def _mock_tensor_from_data(data: Any) -> MagicMock: 

126 """Return a MagicMock representing a tensor from given data.""" 

127 m = MagicMock() 

128 if hasattr(data, "__len__"): 

129 m.shape = (len(data),) 

130 else: 

131 m.shape = () 

132 m.data = data 

133 m.ndim = len(m.shape) 

134 return m 

135 

136 def get_sub_modules(self) -> Dict[str, Any]: 

137 return { 

138 "torch": self, 

139 "torch.fft": MagicMock(), 

140 "torch.nn": MagicMock(), 

141 "torch.nn.functional": MagicMock(), 

142 # "torch.nn": self.nn, 

143 # "torch.nn.Module": FakeModule, 

144 # "torch.nn.LayerNorm": FakeLayerNorm, 

145 # "torch.nn.functional": self.nn.functional, 

146 "torch.nn.parallel": MagicMock(), 

147 "torch.nn.parallel.distributed": MagicMock(), 

148 "torch.nn.modules": MagicMock(), 

149 "torch.nn.modules.utils": MagicMock(), 

150 "torch.nn.common_types": MagicMock(), 

151 "torch.amp": MagicMock(), 

152 "torch.amp.grad_scaler": MagicMock(), 

153 "torch.distributed": MagicMock(), 

154 "torch.distributed.rpc": MagicMock(), 

155 "torch.distributed.algorithms": MagicMock(), 

156 "torch.distributed.algorithms.join": MagicMock(), 

157 "torch.optim": MagicMock(), 

158 "torch.optim.lr_scheduler": MagicMock(), 

159 "torch.utils": MagicMock(), 

160 "torch.utils.data": MagicMock(), 

161 "torch.utils.hooks": MagicMock(), 

162 "torch.utils.checkpoint": MagicMock(), 

163 "torch.utils.serialization": MagicMock(), 

164 "torch.utils.model_zoo": MagicMock(), 

165 "torch.utils._ordered_set": MagicMock(), 

166 "torch.utils._sympy": MagicMock(), 

167 "torch.utils._sympy.functions": MagicMock(), 

168 "torch.utils._pytree": MagicMock(), 

169 "torch.cuda": MagicMock(), 

170 }