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

1# Copyright 2023 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"""Overload of numpy.iinfo to handle dtypes defined in ml_dtypes.""" 

16 

17from ml_dtypes._ml_dtypes_ext import int1 

18from ml_dtypes._ml_dtypes_ext import int2 

19from ml_dtypes._ml_dtypes_ext import int4 

20from ml_dtypes._ml_dtypes_ext import uint1 

21from ml_dtypes._ml_dtypes_ext import uint2 

22from ml_dtypes._ml_dtypes_ext import uint4 

23import numpy as np 

24 

25_int1_dtype = np.dtype(int1) 

26_uint1_dtype = np.dtype(uint1) 

27_int2_dtype = np.dtype(int2) 

28_uint2_dtype = np.dtype(uint2) 

29_int4_dtype = np.dtype(int4) 

30_uint4_dtype = np.dtype(uint4) 

31 

32 

33class iinfo: # pylint: disable=invalid-name,missing-class-docstring 

34 kind: str 

35 bits: int 

36 min: int 

37 max: int 

38 dtype: np.dtype 

39 

40 def __init__(self, int_type): 

41 # Check for dtype attribute in order to handle finfo(arr), as required by 

42 # the Python Array API standard. 

43 if hasattr(int_type, "dtype") and isinstance(int_type.dtype, np.dtype): 

44 int_type = int_type.dtype 

45 

46 if int_type == _int1_dtype: 

47 self.dtype = _int1_dtype 

48 self.kind = "i" 

49 self.bits = 1 

50 self.min = -1 

51 self.max = 0 

52 elif int_type == _uint1_dtype: 

53 self.dtype = _uint1_dtype 

54 self.kind = "u" 

55 self.bits = 1 

56 self.min = 0 

57 self.max = 1 

58 elif int_type == _int2_dtype: 

59 self.dtype = _int2_dtype 

60 self.kind = "i" 

61 self.bits = 2 

62 self.min = -2 

63 self.max = 1 

64 elif int_type == _uint2_dtype: 

65 self.dtype = _uint2_dtype 

66 self.kind = "u" 

67 self.bits = 2 

68 self.min = 0 

69 self.max = 3 

70 elif int_type == _int4_dtype: 

71 self.dtype = _int4_dtype 

72 self.kind = "i" 

73 self.bits = 4 

74 self.min = -8 

75 self.max = 7 

76 elif int_type == _uint4_dtype: 

77 self.dtype = _uint4_dtype 

78 self.kind = "u" 

79 self.bits = 4 

80 self.min = 0 

81 self.max = 15 

82 else: 

83 ii = np.iinfo(int_type) 

84 self.dtype = ii.dtype 

85 self.kind = ii.kind 

86 self.bits = ii.bits 

87 self.min = ii.min 

88 self.max = ii.max 

89 

90 def __repr__(self): 

91 return f"iinfo(min={self.min}, max={self.max}, dtype={self.dtype})" 

92 

93 def __str__(self): 

94 return repr(self)