|
5 | 5 | from __future__ import annotations |
6 | 6 |
|
7 | 7 | import inspect |
8 | | -from typing import NamedTuple, Optional, Sequence, Tuple, Union |
| 8 | +from typing import Any, NamedTuple, Optional, Sequence, Tuple, Union |
9 | 9 |
|
10 | 10 | from ._typing import Array, Device, DType, Namespace |
11 | 11 | from ._helpers import ( |
@@ -609,13 +609,30 @@ def sign(x: Array, /, xp: Namespace, **kwargs) -> Array: |
609 | 609 | out[xp.isnan(x)] = xp.nan |
610 | 610 | return out[()] |
611 | 611 |
|
| 612 | + |
| 613 | +def finfo(type_: DType | Array, /, xp: Namespace) -> Any: |
| 614 | + # It is surprisingly difficult to recognize a dtype apart from an array. |
| 615 | + # np.int64 is not the same as np.asarray(1).dtype! |
| 616 | + try: |
| 617 | + return xp.finfo(type_) |
| 618 | + except (ValueError, TypeError): |
| 619 | + return xp.finfo(type_.dtype) |
| 620 | + |
| 621 | + |
| 622 | +def iinfo(type_: DType | Array, /, xp: Namespace) -> Any: |
| 623 | + try: |
| 624 | + return xp.iinfo(type_) |
| 625 | + except (ValueError, TypeError): |
| 626 | + return xp.iinfo(type_.dtype) |
| 627 | + |
| 628 | + |
612 | 629 | __all__ = ['arange', 'empty', 'empty_like', 'eye', 'full', 'full_like', |
613 | 630 | 'linspace', 'ones', 'ones_like', 'zeros', 'zeros_like', |
614 | 631 | 'UniqueAllResult', 'UniqueCountsResult', 'UniqueInverseResult', |
615 | 632 | 'unique_all', 'unique_counts', 'unique_inverse', 'unique_values', |
616 | 633 | 'std', 'var', 'cumulative_sum', 'cumulative_prod','clip', 'permute_dims', |
617 | 634 | 'reshape', 'argsort', 'sort', 'nonzero', 'ceil', 'floor', 'trunc', |
618 | 635 | 'matmul', 'matrix_transpose', 'tensordot', 'vecdot', 'isdtype', |
619 | | - 'unstack', 'sign'] |
| 636 | + 'unstack', 'sign', 'finfo', 'iinfo'] |
620 | 637 |
|
621 | 638 | _all_ignore = ['inspect', 'array_namespace', 'NamedTuple'] |
0 commit comments