Coverage for src/flag_gems/runtime/backend/_mthreads/ops/__init__.py: 0%

1from torch_musa import current_device, get_device_capability 

2 

3from .all import all, all_dim, all_dims 

4from .amax import amax 

5from .any import any, any_dim, any_dims 

6from .arange import arange, arange_start 

7from .argmin import argmin 

8from .batch_norm import batch_norm, batch_norm_backward 

9from .celu import celu 

10from .conv2d import conv2d 

11from .dropout import dropout, dropout_backward 

12from .gather import gather, gather_backward 

13from .index_add import index_add, index_add_ 

14from .index_put import _index_put_impl_, index_put, index_put_ 

15from .index_select import index_select 

16from .log import log 

17from .log_softmax import ( 

18 log_softmax, 

19 log_softmax_backward, 

20 log_softmax_backward_out, 

21 log_softmax_out, 

22) 

23from .max import max, max_dim 

24from .min import min, min_dim 

25from .normal import normal_ 

26from .one_hot import one_hot 

27from .ones import ones 

28from .ones_like import ones_like 

29from .prod import prod, prod_dim 

30from .rand import rand 

31from .rand_like import rand_like 

32from .randn import randn 

33from .randn_like import randn_like 

34from .randperm import randperm 

35from .repeat import repeat 

36from .repeat_interleave import ( 

37 repeat_interleave_self_int, 

38 repeat_interleave_self_tensor, 

39 repeat_interleave_tensor, 

40) 

41from .resolve_conj import resolve_conj 

42from .sort import sort, sort_stable 

43from .tile import tile 

44from .unique import _unique2 

45from .w8a8_block_fp8_matmul import w8a8_block_fp8_matmul 

46from .zeros import zero_, zeros 

47from .zeros_like import zeros_like 

48 

49__all__ = [ 

50 "amax", 

51 "all", 

52 "all_dim", 

53 "all_dims", 

54 "any", 

55 "any_dim", 

56 "any_dims", 

57 "arange", 

58 "arange_start", 

59 "argmin", 

60 "batch_norm", 

61 "batch_norm_backward", 

62 "celu", 

63 # "celu_", 

64 "conv2d", 

65 "dropout", 

66 "dropout_backward", 

67 "gather", 

68 "gather_backward", 

69 "index_add", 

70 "index_add_", 

71 "index_put", 

72 "index_put_", 

73 "_index_put_impl_", 

74 "index_select", 

75 "log", 

76 "log_softmax", 

77 "log_softmax_backward", 

78 "log_softmax_backward_out", 

79 "log_softmax_out", 

80 "max", 

81 "max_dim", 

82 "min", 

83 "min_dim", 

84 "normal_", 

85 "one_hot", 

86 "ones", 

87 "ones_like", 

88 "prod", 

89 "prod_dim", 

90 "rand", 

91 "rand_like", 

92 "randn", 

93 "randn_like", 

94 "randperm", 

95 "repeat", 

96 "repeat_interleave_self_int", 

97 "repeat_interleave_self_tensor", 

98 "repeat_interleave_tensor", 

99 "resolve_conj", 

100 "sort", 

101 "sort_stable", 

102 "tile", 

103 "_unique2", 

104 "w8a8_block_fp8_matmul", 

105 "zero_", 

106 "zeros", 

107 "zeros_like", 

108] 

109 

110 

111if get_device_capability(current_device())[0] >= 3: 

112 from .addmm import addmm, addmm_dtype, addmm_dtype_out # noqa: F401 

113 from .bmm import bmm # noqa: F401 

114 from .gelu import gelu # noqa: F401 

115 from .mm import mm # noqa: F401 

116 from .tanh import tanh # noqa: F401 

117 

118 __all__.extend( 

119 [ 

120 "addmm", 

121 "addmm_dtype", 

122 "addmm_dtype_out", 

123 "bmm", 

124 "gelu", 

125 "mm", 

126 "tanh", 

127 ] 

128 )