_array_api_info.py 10 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346
  1. """
  2. Array API Inspection namespace
  3. This is the namespace for inspection functions as defined by the array API
  4. standard. See
  5. https://data-apis.org/array-api/latest/API_specification/inspection.html for
  6. more details.
  7. """
  8. from numpy._core import (
  9. bool,
  10. complex64,
  11. complex128,
  12. dtype,
  13. float32,
  14. float64,
  15. int8,
  16. int16,
  17. int32,
  18. int64,
  19. intp,
  20. uint8,
  21. uint16,
  22. uint32,
  23. uint64,
  24. )
  25. class __array_namespace_info__:
  26. """
  27. Get the array API inspection namespace for NumPy.
  28. The array API inspection namespace defines the following functions:
  29. - capabilities()
  30. - default_device()
  31. - default_dtypes()
  32. - dtypes()
  33. - devices()
  34. See
  35. https://data-apis.org/array-api/latest/API_specification/inspection.html
  36. for more details.
  37. Returns
  38. -------
  39. info : ModuleType
  40. The array API inspection namespace for NumPy.
  41. Examples
  42. --------
  43. >>> info = np.__array_namespace_info__()
  44. >>> info.default_dtypes()
  45. {'real floating': numpy.float64,
  46. 'complex floating': numpy.complex128,
  47. 'integral': numpy.int64,
  48. 'indexing': numpy.int64}
  49. """
  50. __module__ = 'numpy'
  51. def capabilities(self):
  52. """
  53. Return a dictionary of array API library capabilities.
  54. The resulting dictionary has the following keys:
  55. - **"boolean indexing"**: boolean indicating whether an array library
  56. supports boolean indexing. Always ``True`` for NumPy.
  57. - **"data-dependent shapes"**: boolean indicating whether an array
  58. library supports data-dependent output shapes. Always ``True`` for
  59. NumPy.
  60. See
  61. https://data-apis.org/array-api/latest/API_specification/generated/array_api.info.capabilities.html
  62. for more details.
  63. See Also
  64. --------
  65. __array_namespace_info__.default_device,
  66. __array_namespace_info__.default_dtypes,
  67. __array_namespace_info__.dtypes,
  68. __array_namespace_info__.devices
  69. Returns
  70. -------
  71. capabilities : dict
  72. A dictionary of array API library capabilities.
  73. Examples
  74. --------
  75. >>> info = np.__array_namespace_info__()
  76. >>> info.capabilities()
  77. {'boolean indexing': True,
  78. 'data-dependent shapes': True,
  79. 'max dimensions': 64}
  80. """
  81. return {
  82. "boolean indexing": True,
  83. "data-dependent shapes": True,
  84. "max dimensions": 64,
  85. }
  86. def default_device(self):
  87. """
  88. The default device used for new NumPy arrays.
  89. For NumPy, this always returns ``'cpu'``.
  90. See Also
  91. --------
  92. __array_namespace_info__.capabilities,
  93. __array_namespace_info__.default_dtypes,
  94. __array_namespace_info__.dtypes,
  95. __array_namespace_info__.devices
  96. Returns
  97. -------
  98. device : str
  99. The default device used for new NumPy arrays.
  100. Examples
  101. --------
  102. >>> info = np.__array_namespace_info__()
  103. >>> info.default_device()
  104. 'cpu'
  105. """
  106. return "cpu"
  107. def default_dtypes(self, *, device=None):
  108. """
  109. The default data types used for new NumPy arrays.
  110. For NumPy, this always returns the following dictionary:
  111. - **"real floating"**: ``numpy.float64``
  112. - **"complex floating"**: ``numpy.complex128``
  113. - **"integral"**: ``numpy.intp``
  114. - **"indexing"**: ``numpy.intp``
  115. Parameters
  116. ----------
  117. device : str, optional
  118. The device to get the default data types for. For NumPy, only
  119. ``'cpu'`` is allowed.
  120. Returns
  121. -------
  122. dtypes : dict
  123. A dictionary describing the default data types used for new NumPy
  124. arrays.
  125. See Also
  126. --------
  127. __array_namespace_info__.capabilities,
  128. __array_namespace_info__.default_device,
  129. __array_namespace_info__.dtypes,
  130. __array_namespace_info__.devices
  131. Examples
  132. --------
  133. >>> info = np.__array_namespace_info__()
  134. >>> info.default_dtypes()
  135. {'real floating': numpy.float64,
  136. 'complex floating': numpy.complex128,
  137. 'integral': numpy.int64,
  138. 'indexing': numpy.int64}
  139. """
  140. if device not in ["cpu", None]:
  141. raise ValueError(
  142. 'Device not understood. Only "cpu" is allowed, but received:'
  143. f' {device}'
  144. )
  145. return {
  146. "real floating": dtype(float64),
  147. "complex floating": dtype(complex128),
  148. "integral": dtype(intp),
  149. "indexing": dtype(intp),
  150. }
  151. def dtypes(self, *, device=None, kind=None):
  152. """
  153. The array API data types supported by NumPy.
  154. Note that this function only returns data types that are defined by
  155. the array API.
  156. Parameters
  157. ----------
  158. device : str, optional
  159. The device to get the data types for. For NumPy, only ``'cpu'`` is
  160. allowed.
  161. kind : str or tuple of str, optional
  162. The kind of data types to return. If ``None``, all data types are
  163. returned. If a string, only data types of that kind are returned.
  164. If a tuple, a dictionary containing the union of the given kinds
  165. is returned. The following kinds are supported:
  166. - ``'bool'``: boolean data types (i.e., ``bool``).
  167. - ``'signed integer'``: signed integer data types (i.e., ``int8``,
  168. ``int16``, ``int32``, ``int64``).
  169. - ``'unsigned integer'``: unsigned integer data types (i.e.,
  170. ``uint8``, ``uint16``, ``uint32``, ``uint64``).
  171. - ``'integral'``: integer data types. Shorthand for ``('signed
  172. integer', 'unsigned integer')``.
  173. - ``'real floating'``: real-valued floating-point data types
  174. (i.e., ``float32``, ``float64``).
  175. - ``'complex floating'``: complex floating-point data types (i.e.,
  176. ``complex64``, ``complex128``).
  177. - ``'numeric'``: numeric data types. Shorthand for ``('integral',
  178. 'real floating', 'complex floating')``.
  179. Returns
  180. -------
  181. dtypes : dict
  182. A dictionary mapping the names of data types to the corresponding
  183. NumPy data types.
  184. See Also
  185. --------
  186. __array_namespace_info__.capabilities,
  187. __array_namespace_info__.default_device,
  188. __array_namespace_info__.default_dtypes,
  189. __array_namespace_info__.devices
  190. Examples
  191. --------
  192. >>> info = np.__array_namespace_info__()
  193. >>> info.dtypes(kind='signed integer')
  194. {'int8': numpy.int8,
  195. 'int16': numpy.int16,
  196. 'int32': numpy.int32,
  197. 'int64': numpy.int64}
  198. """
  199. if device not in ["cpu", None]:
  200. raise ValueError(
  201. 'Device not understood. Only "cpu" is allowed, but received:'
  202. f' {device}'
  203. )
  204. if kind is None:
  205. return {
  206. "bool": dtype(bool),
  207. "int8": dtype(int8),
  208. "int16": dtype(int16),
  209. "int32": dtype(int32),
  210. "int64": dtype(int64),
  211. "uint8": dtype(uint8),
  212. "uint16": dtype(uint16),
  213. "uint32": dtype(uint32),
  214. "uint64": dtype(uint64),
  215. "float32": dtype(float32),
  216. "float64": dtype(float64),
  217. "complex64": dtype(complex64),
  218. "complex128": dtype(complex128),
  219. }
  220. if kind == "bool":
  221. return {"bool": bool}
  222. if kind == "signed integer":
  223. return {
  224. "int8": dtype(int8),
  225. "int16": dtype(int16),
  226. "int32": dtype(int32),
  227. "int64": dtype(int64),
  228. }
  229. if kind == "unsigned integer":
  230. return {
  231. "uint8": dtype(uint8),
  232. "uint16": dtype(uint16),
  233. "uint32": dtype(uint32),
  234. "uint64": dtype(uint64),
  235. }
  236. if kind == "integral":
  237. return {
  238. "int8": dtype(int8),
  239. "int16": dtype(int16),
  240. "int32": dtype(int32),
  241. "int64": dtype(int64),
  242. "uint8": dtype(uint8),
  243. "uint16": dtype(uint16),
  244. "uint32": dtype(uint32),
  245. "uint64": dtype(uint64),
  246. }
  247. if kind == "real floating":
  248. return {
  249. "float32": dtype(float32),
  250. "float64": dtype(float64),
  251. }
  252. if kind == "complex floating":
  253. return {
  254. "complex64": dtype(complex64),
  255. "complex128": dtype(complex128),
  256. }
  257. if kind == "numeric":
  258. return {
  259. "int8": dtype(int8),
  260. "int16": dtype(int16),
  261. "int32": dtype(int32),
  262. "int64": dtype(int64),
  263. "uint8": dtype(uint8),
  264. "uint16": dtype(uint16),
  265. "uint32": dtype(uint32),
  266. "uint64": dtype(uint64),
  267. "float32": dtype(float32),
  268. "float64": dtype(float64),
  269. "complex64": dtype(complex64),
  270. "complex128": dtype(complex128),
  271. }
  272. if isinstance(kind, tuple):
  273. res = {}
  274. for k in kind:
  275. res.update(self.dtypes(kind=k))
  276. return res
  277. raise ValueError(f"unsupported kind: {kind!r}")
  278. def devices(self):
  279. """
  280. The devices supported by NumPy.
  281. For NumPy, this always returns ``['cpu']``.
  282. Returns
  283. -------
  284. devices : list of str
  285. The devices supported by NumPy.
  286. See Also
  287. --------
  288. __array_namespace_info__.capabilities,
  289. __array_namespace_info__.default_device,
  290. __array_namespace_info__.default_dtypes,
  291. __array_namespace_info__.dtypes
  292. Examples
  293. --------
  294. >>> info = np.__array_namespace_info__()
  295. >>> info.devices()
  296. ['cpu']
  297. """
  298. return ["cpu"]