Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/ml_dtypes/__init__.py: 78%

1# Copyright 2022 The ml_dtypes Authors. 

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__version__ = "0.6.0" 

16__all__ = [ 

17 "__version__", 

18 "bfloat16", 

19 "finfo", 

20 "bcomplex32", 

21 "complex32", 

22 "float4_e2m1fn", 

23 "float6_e2m3fn", 

24 "float6_e3m2fn", 

25 "float8_e3m4", 

26 "float8_e4m3", 

27 "float8_e4m3b11fnuz", 

28 "float8_e4m3fn", 

29 "float8_e4m3fnuz", 

30 "float8_e5m2", 

31 "float8_e5m2fnuz", 

32 "float8_e8m0fnu", 

33 "iinfo", 

34 "int1", 

35 "int2", 

36 "int4", 

37 "uint1", 

38 "uint2", 

39 "uint4", 

40] 

41 

42from ml_dtypes._finfo import finfo 

43from ml_dtypes._iinfo import iinfo 

44from ml_dtypes._ml_dtypes_ext import bcomplex32 

45from ml_dtypes._ml_dtypes_ext import bfloat16 

46from ml_dtypes._ml_dtypes_ext import complex32 

47from ml_dtypes._ml_dtypes_ext import float4_e2m1fn 

48from ml_dtypes._ml_dtypes_ext import float6_e2m3fn 

49from ml_dtypes._ml_dtypes_ext import float6_e3m2fn 

50from ml_dtypes._ml_dtypes_ext import float8_e3m4 

51from ml_dtypes._ml_dtypes_ext import float8_e4m3 

52from ml_dtypes._ml_dtypes_ext import float8_e4m3b11fnuz 

53from ml_dtypes._ml_dtypes_ext import float8_e4m3fn 

54from ml_dtypes._ml_dtypes_ext import float8_e4m3fnuz 

55from ml_dtypes._ml_dtypes_ext import float8_e5m2 

56from ml_dtypes._ml_dtypes_ext import float8_e5m2fnuz 

57from ml_dtypes._ml_dtypes_ext import float8_e8m0fnu 

58from ml_dtypes._ml_dtypes_ext import int1 

59from ml_dtypes._ml_dtypes_ext import int2 

60from ml_dtypes._ml_dtypes_ext import int4 

61from ml_dtypes._ml_dtypes_ext import uint1 

62from ml_dtypes._ml_dtypes_ext import uint2 

63from ml_dtypes._ml_dtypes_ext import uint4 

64import numpy as _np 

65 

66bfloat16: type[_np.generic] 

67float4_e2m1fn: type[_np.generic] 

68float6_e2m3fn: type[_np.generic] 

69float6_e3m2fn: type[_np.generic] 

70float8_e3m4: type[_np.generic] 

71float8_e4m3: type[_np.generic] 

72float8_e4m3b11fnuz: type[_np.generic] 

73float8_e4m3fn: type[_np.generic] 

74float8_e4m3fnuz: type[_np.generic] 

75float8_e5m2: type[_np.generic] 

76float8_e5m2fnuz: type[_np.generic] 

77float8_e8m0fnu: type[_np.generic] 

78int1: type[_np.generic] 

79int2: type[_np.generic] 

80int4: type[_np.generic] 

81uint1: type[_np.generic] 

82uint2: type[_np.generic] 

83uint4: type[_np.generic] 

84bcomplex32: type[_np.generic] 

85complex32: type[_np.generic] 

86 

87 

88def real(x: _np.ndarray) -> _np.ndarray: 

89 """Return the real part of a complex array. 

90 

91 This is a helper that uses `x.real` except for NumPy arrays of 

92 bcomplex32 or complex32. NumPy cannot correctly understand that these 

93 are complex dtypes as of NumPy 2.4 at least. 

94 

95 Args: 

96 x: The input array. 

97 

98 Returns: 

99 The real part of the input array. 

100 """ 

101 if isinstance(x, _np.ndarray): 

102 # Use a view. We add an axes to ensure it is contiguous. 

103 if x.dtype.type is bcomplex32: 

104 return x[..., None].view(bfloat16)[..., 0] 

105 elif x.dtype.type is complex32: 

106 return x[..., None].view(_np.float16)[..., 0] 

107 

108 # Otherwise, assume everything is OK with just using the normal `.real` 

109 return x.real 

110 

111 

112def imag(x: _np.ndarray) -> _np.ndarray: 

113 """Return the imaginary part of a complex array. 

114 

115 This is a helper that uses `x.imag` except for NumPy arrays of 

116 bcomplex32 or complex32. NumPy cannot correctly understand that these 

117 are complex dtypes as of NumPy 2.4 at least. 

118 

119 Args: 

120 x: The input array. 

121 

122 Returns: 

123 The imaginary part of the input array. 

124 """ 

125 if isinstance(x, _np.ndarray): 

126 # Use a view. We add an axes to ensure it is contiguous. 

127 if x.dtype.type is bcomplex32: 

128 return x[..., None].view(bfloat16)[..., 1] 

129 elif x.dtype.type is complex32: 

130 return x[..., None].view(_np.float16)[..., 1] 

131 

132 # Otherwise, assume everything is OK with just using the normal `.imag` 

133 return x.imag