xbarray 0.0.1a1__py3-none-any.whl

xbarray/numpy/_extra.py

@@ -0,0 +1,83 @@
+from typing import Any, Union, Optional
+import numpy as np
+from ._typing import ARRAY_TYPE, DTYPE_TYPE, DEVICE_TYPE, RNG_TYPE
+from xbarray.base import ComputeBackend, SupportsDLPack
+
+__all__ = [
+    "default_integer_dtype",
+    "default_floating_dtype",
+    "default_boolean_dtype",
+    "is_backendarray",
+    "from_numpy",
+    "from_other_backend",
+    "to_numpy",
+    "to_dlpack",
+    "dtype_is_real_integer",
+    "dtype_is_real_floating",
+    "dtype_is_boolean",
+    "abbreviate_array",
+]
+
+default_integer_dtype = int
+default_floating_dtype = float
+default_boolean_dtype = bool
+
+def is_backendarray(data : Any) -> bool:
+    return isinstance(data, np.ndarray)
+
+def from_numpy(
+    data : np.ndarray,
+    /,
+    *,
+    dtype : Optional[DTYPE_TYPE] = None,
+    device : Optional[DEVICE_TYPE] = None
+) -> ARRAY_TYPE:
+    return data
+
+def from_other_backend(
+    other_backend: ComputeBackend,
+    data: Any,
+    /,
+) -> ARRAY_TYPE:
+    return other_backend.to_numpy(data)
+
+def to_numpy(
+    data : ARRAY_TYPE
+) -> np.ndarray:
+    return data
+
+def to_dlpack(
+    self,
+    data: ARRAY_TYPE,
+    /,
+) -> SupportsDLPack:
+    return data
+
+def dtype_is_real_integer(
+    dtype: DTYPE_TYPE
+) -> bool:
+    return np.issubdtype(dtype, np.integer)
+
+def dtype_is_real_floating(
+    dtype: DTYPE_TYPE
+) -> bool:
+    return np.issubdtype(dtype, np.floating)
+
+def dtype_is_boolean(
+    dtype: DTYPE_TYPE
+) -> bool:
+    return dtype == np.bool_ or dtype == bool
+
+from xbarray.common.implementations import get_abbreviate_array_function, get_map_fn_over_arrays_function
+from array_api_compat import numpy as compat_module
+abbreviate_array = get_abbreviate_array_function(
+    backend=compat_module,
+    default_integer_dtype=default_integer_dtype,
+    func_dtype_is_real_floating=dtype_is_real_floating,
+    func_dtype_is_real_integer=dtype_is_real_integer,
+    func_dtype_is_boolean=dtype_is_boolean,
+)
+
+map_fn_over_arrays = get_map_fn_over_arrays_function(
+    is_backendarray=is_backendarray,
+)

xbarray/numpy/_typing.py

@@ -0,0 +1,14 @@
+from typing import Union, Any, Optional
+import numpy as np
+
+__all__ = [
+    'ARRAY_TYPE',
+    'DTYPE_TYPE',
+    'DEVICE_TYPE',
+    'RNG_TYPE',
+]
+
+ARRAY_TYPE = np.ndarray
+DTYPE_TYPE = np.dtype
+DEVICE_TYPE = Any
+RNG_TYPE = np.random.Generator

xbarray/numpy/random.py

@@ -0,0 +1,106 @@
+from typing import Union, Optional, Tuple, Any
+from ._typing import ARRAY_TYPE, DTYPE_TYPE, DEVICE_TYPE, RNG_TYPE
+import numpy as np
+
+__all__ = [
+    "random_number_generator",
+    "random_discrete_uniform",
+    "random_uniform",
+    "random_exponential",
+    "random_normal",
+    "random_geometric",
+    "random_permutation"
+]
+
+def random_number_generator(
+    seed : Optional[int] = None,
+    *,
+    device : Optional[DEVICE_TYPE] = None
+) -> RNG_TYPE:
+    return np.random.default_rng(seed)
+
+def random_discrete_uniform(
+    shape : Union[int, Tuple[int, ...]], 
+    from_num : int, 
+    to_num : int, 
+    /,
+    *,
+    rng : RNG_TYPE, 
+    dtype : Optional[DTYPE_TYPE] = None, 
+    device : Optional[DEVICE_TYPE] = None
+) -> Tuple[RNG_TYPE, ARRAY_TYPE]:
+    t = rng.integers(int(from_num), int(to_num), size=shape)
+    if dtype is not None:
+        t = t.astype(dtype)
+    return rng, t
+
+def random_uniform(
+    shape: Union[int, Tuple[int, ...]], 
+    /,
+    *,
+    rng : RNG_TYPE, 
+    low : float = 0.0, high : float = 1.0,
+    dtype : Optional[DTYPE_TYPE] = None, 
+    device : Optional[DEVICE_TYPE] = None
+) -> Tuple[RNG_TYPE, ARRAY_TYPE]:
+    t = rng.uniform(float(low), float(high), size=shape)
+    if dtype is not None:
+        t = t.astype(dtype)
+    return rng, t
+
+def random_exponential(
+    shape: Union[int, Tuple[int, ...]], 
+    /,
+    *,
+    rng : RNG_TYPE,  
+    lambd : float = 1.0, 
+    dtype : Optional[DTYPE_TYPE] = None, 
+    device : Optional[DEVICE_TYPE] = None
+) -> Tuple[RNG_TYPE, ARRAY_TYPE]:
+    t = rng.exponential(1.0 / float(lambd), size=shape)
+    if dtype is not None:
+        t = t.astype(dtype)
+    return rng, t
+
+@classmethod
+def random_normal(
+    shape: Union[int, Tuple[int, ...]],
+    /,
+    *, 
+    rng : RNG_TYPE,
+    mean : float = 0.0, std : float = 1.0, 
+    dtype : Optional[DTYPE_TYPE] = None, 
+    device : Optional[Any] = None
+) -> Tuple[RNG_TYPE, ARRAY_TYPE]:
+    t = rng.normal(mean, std, size=shape)
+    if dtype is not None:
+        t = t.astype(dtype)
+    return rng, t
+
+
+def random_geometric(
+    shape: Union[int, Tuple[int, ...]], 
+    /,
+    *,
+    p: float, 
+    rng: RNG_TYPE,
+    dtype: Optional[DTYPE_TYPE] = None, 
+    device: Optional[Any] = None
+) -> Tuple[RNG_TYPE, ARRAY_TYPE]:
+    t = rng.geometric(p, size=shape)
+    if dtype is not None:
+        t = t.astype(dtype)
+    return rng, t
+
+def random_permutation(
+    n : int,
+    /,
+    *,
+    rng: RNG_TYPE,
+    dtype: Optional[DTYPE_TYPE] = None,
+    device: Optional[DEVICE_TYPE] = None
+) -> Tuple[RNG_TYPE, ARRAY_TYPE]:
+    t = rng.permutation(n)
+    if dtype is not None:
+        t = t.astype(dtype)
+    return rng, t

xbarray/pytorch/__init__.py

@@ -0,0 +1,20 @@
+from array_api_compat.torch import *
+
+simplified_name = "pytorch"
+
+from array_api_compat import torch as compat_module
+
+# Import and bind all functions from array_api_extra before exposing them
+import array_api_extra
+from functools import partial
+for api_name in dir(array_api_extra):
+    if api_name.startswith('_'):
+        continue
+    globals()[api_name] = partial(
+        getattr(array_api_extra, api_name),
+        xp=compat_module
+    )
+
+from ._typing import *
+from ._extra import *
+__import__(__package__ + ".random")

xbarray/pytorch/_extra.py

@@ -0,0 +1,109 @@
+from typing import Any, Union, Optional
+import numpy as np
+import torch
+from ._typing import ARRAY_TYPE, DTYPE_TYPE, DEVICE_TYPE, RNG_TYPE
+from xbarray.base import ComputeBackend, SupportsDLPack
+
+PYTORCH_DTYPE_CAST_MAP = {
+    torch.uint16: torch.int16,
+    torch.uint32: torch.int32,
+    torch.uint64: torch.int64,
+    torch.float8_e4m3fn: torch.float16,
+    torch.float8_e5m2: torch.float16,
+}
+
+__all__ = [
+    "default_integer_dtype",
+    "default_floating_dtype",
+    "default_boolean_dtype",
+    "is_backendarray",
+    "from_numpy",
+    "from_other_backend",
+    "to_numpy",
+    "to_dlpack",
+    "dtype_is_real_integer",
+    "dtype_is_real_floating",
+    "dtype_is_boolean",
+    "abbreviate_array",
+]
+
+default_integer_dtype = torch.int32
+default_floating_dtype = torch.float32
+default_boolean_dtype = torch.bool
+
+def is_backendarray(data : Any) -> bool:
+    return isinstance(data, torch.Tensor)
+
+def from_numpy(
+    data : np.ndarray,
+    /,
+    *,
+    dtype : Optional[DTYPE_TYPE] = None,
+    device : Optional[DEVICE_TYPE] = None
+) -> ARRAY_TYPE:
+    t = torch.from_numpy(data)
+    target_dtype = dtype if dtype is not None else PYTORCH_DTYPE_CAST_MAP.get(t.dtype, t.dtype)
+    if target_dtype is not None or device is not None:
+        t = t.to(device=device, dtype=target_dtype)
+    return t
+
+def from_other_backend(
+    other_backend: ComputeBackend,
+    data: Any,
+    /,
+) -> ARRAY_TYPE:
+    dat_dlpack = other_backend.to_dlpack(data)
+    return torch.from_dlpack(dat_dlpack)
+
+def to_numpy(
+    data : ARRAY_TYPE
+) -> np.ndarray:
+    # Torch bfloat16 is not supported by numpy
+    if data.dtype == torch.bfloat16:
+        data = data.to(torch.float32)
+    return data.cpu().numpy()
+
+def to_dlpack(
+    self,
+    data: ARRAY_TYPE,
+    /,
+) -> SupportsDLPack:
+    return data
+
+def dtype_is_real_integer(
+    dtype: DTYPE_TYPE
+) -> bool:
+    # https://pytorch.org/docs/stable/tensors.html#id12
+    return dtype in [
+        torch.int8, torch.int16, torch.int32, torch.int64,
+        torch.uint8, 
+        torch.int,
+        torch.long
+    ]
+
+def dtype_is_real_floating(
+    dtype: DTYPE_TYPE
+) -> bool:
+    return dtype in [
+        torch.float16, torch.float32, torch.float64, 
+        torch.float, torch.double, 
+        torch.bfloat16
+    ]
+
+def dtype_is_boolean(
+    dtype: DTYPE_TYPE
+) -> bool:
+    return dtype == torch.bool
+
+from xbarray.common.implementations import get_abbreviate_array_function, get_map_fn_over_arrays_function
+from array_api_compat import torch as compat_module
+abbreviate_array = get_abbreviate_array_function(
+    compat_module, 
+    default_integer_dtype=default_integer_dtype, 
+    func_dtype_is_real_floating=dtype_is_real_floating,
+    func_dtype_is_real_integer=dtype_is_real_integer,
+    func_dtype_is_boolean=dtype_is_boolean
+)
+map_fn_over_arrays = get_map_fn_over_arrays_function(
+    is_backendarray=is_backendarray,
+)

xbarray/pytorch/_typing.py

@@ -0,0 +1,13 @@
+import torch
+
+__all__ = [
+    'ARRAY_TYPE',
+    'DTYPE_TYPE',
+    'DEVICE_TYPE',
+    'RNG_TYPE',
+]
+
+ARRAY_TYPE = torch.Tensor
+DTYPE_TYPE = torch.dtype
+DEVICE_TYPE = torch.device
+RNG_TYPE = torch.Generator

xbarray/pytorch/random.py

@@ -0,0 +1,102 @@
+from typing import Union, Optional, Tuple, Any
+from ._typing import ARRAY_TYPE, DTYPE_TYPE, DEVICE_TYPE, RNG_TYPE
+import torch
+
+__all__ = [
+    "random_number_generator",
+    "random_discrete_uniform",
+    "random_uniform",
+    "random_exponential",
+    "random_normal",
+    "random_geometric",
+    "random_permutation"
+]
+
+def random_number_generator(
+    seed : Optional[int] = None,
+    *,
+    device : Optional[DEVICE_TYPE] = None
+) -> RNG_TYPE:
+    rng = torch.Generator(
+        device=device
+    )
+    if seed is not None:
+        rng = rng.manual_seed(seed)
+    return rng
+
+
+def random_discrete_uniform(
+    shape : Union[int, Tuple[int, ...]], 
+    from_num : int, 
+    to_num : int, 
+    /,
+    *,
+    rng : RNG_TYPE, 
+    dtype : Optional[DTYPE_TYPE] = None, 
+    device : Optional[DEVICE_TYPE] = None
+) -> Tuple[RNG_TYPE, ARRAY_TYPE]:
+    t = torch.randint(int(from_num), int(to_num), shape, generator=rng, dtype=dtype, device=device)
+    return rng, t
+
+def random_uniform(
+    shape: Union[int, Tuple[int, ...]], 
+    /,
+    *,
+    rng : RNG_TYPE, 
+    low : float = 0.0, high : float = 1.0,
+    dtype : Optional[DTYPE_TYPE] = None, 
+    device : Optional[DEVICE_TYPE] = None
+) -> Tuple[RNG_TYPE, ARRAY_TYPE]:
+    t = torch.rand(shape, generator=rng, dtype=dtype, device=device)
+    t = t * (high - low) + low
+    return rng, t
+
+def random_exponential(
+    shape: Union[int, Tuple[int, ...]], 
+    /,
+    *,
+    rng : RNG_TYPE,  
+    lambd : float = 1.0, 
+    dtype : Optional[DTYPE_TYPE] = None, 
+    device : Optional[DEVICE_TYPE] = None
+) -> Tuple[RNG_TYPE, ARRAY_TYPE]:
+    t = torch.empty(shape, dtype=dtype, device=device)
+    t = t.exponential_(lambd, generator=rng)
+    return rng, t
+
+@classmethod
+def random_normal(
+    shape: Union[int, Tuple[int, ...]],
+    /,
+    *, 
+    rng : RNG_TYPE,
+    mean : float = 0.0, std : float = 1.0, 
+    dtype : Optional[DTYPE_TYPE] = None, 
+    device : Optional[Any] = None
+) -> Tuple[RNG_TYPE, ARRAY_TYPE]:
+    t = torch.normal(mean, std, shape, generator=rng, dtype=dtype, device=device)
+    return rng, t
+
+def random_geometric(
+    shape: Union[int, Tuple[int, ...]], 
+    /,
+    *,
+    p: float, 
+    rng: RNG_TYPE,
+    dtype: Optional[DTYPE_TYPE] = None, 
+    device: Optional[Any] = None
+) -> Tuple[RNG_TYPE, ARRAY_TYPE]:
+    t = torch.empty(shape, dtype=dtype, device=device)
+    t = t.geometric_(p, generator=rng)
+    return rng, t
+
+def random_permutation(
+    n : int,
+    /,
+    *,
+    rng: RNG_TYPE,
+    dtype: Optional[DTYPE_TYPE] = None,
+    device: Optional[DEVICE_TYPE] = None
+) -> Tuple[RNG_TYPE, ARRAY_TYPE]:
+    t = torch.randperm(n, generator=rng, dtype=dtype, device=device)
+    return rng, t

xbarray-0.0.1a1.dist-info/METADATA

@@ -0,0 +1,14 @@
+Metadata-Version: 2.4
+Name: xbarray
+Version: 0.0.1a1
+Requires-Python: >=3.10
+License-File: LICENSE
+Requires-Dist: typing_extensions>=4.5
+Requires-Dist: array_api_compat
+Requires-Dist: array_api_extra
+Requires-Dist: numpy
+Provides-Extra: torch
+Requires-Dist: torch; extra == "torch"
+Provides-Extra: jax
+Requires-Dist: jax; extra == "jax"
+Dynamic: license-file

xbarray-0.0.1a1.dist-info/RECORD

@@ -0,0 +1,35 @@
+array_api_typing/__init__.py,sha256=5vcE63PsZ_utc7j2VmLK5evUhar-942p95HyjTiLOc0,179
+array_api_typing/typing_2024_12/__init__.py,sha256=JuD2yojHl3eQI74U75tFhSX5rSy32TR5nKyi714ZfX4,347
+array_api_typing/typing_2024_12/_api_constant.py,sha256=tahdzdi7vAZ3UnM8oDcUqvx2FFsMTAyUbuwgarU6K9U,676
+array_api_typing/typing_2024_12/_api_fft_typing.py,sha256=i48S-9trw3s72xiCJmeCNGlJLwwOVq0kJC4lD9se_co,38641
+array_api_typing/typing_2024_12/_api_linalg_typing.py,sha256=Dt5fDTwYb-GAmP7duajwPqH3zaJmwvxDEOOzdl5X0LE,48950
+array_api_typing/typing_2024_12/_api_return_typing.py,sha256=7TxSBRv0H0uQ9VnixbyuEXJQnl1uESo102RsAloI9SA,2302
+array_api_typing/typing_2024_12/_api_typing.py,sha256=rFuF7E6b2RrijzQBrZx6_dYEn9Yslq8FqI31L1Ccrf4,294255
+array_api_typing/typing_2024_12/_array_typing.py,sha256=GMqs5LgshujRIh7i0mGtZy4Vt56FSXqBmNlBGidFRnc,57403
+array_api_typing/typing_compat/__init__.py,sha256=JuD2yojHl3eQI74U75tFhSX5rSy32TR5nKyi714ZfX4,347
+array_api_typing/typing_compat/_api_typing.py,sha256=RnGZFD8AEIqOCurcI68xyw7y9W9hm4yWUGf7NlJkNk0,930
+array_api_typing/typing_compat/_array_typing.py,sha256=CVv91wDjjG-8kSofkqgA85qpM-3x34Zh_SNuBFegW9o,1182
+array_api_typing/typing_extra/__init__.py,sha256=YfdhD-Sfk3SCfI9lHmA-PbVLzms1OFF5x0ekzI3aafk,323
+array_api_typing/typing_extra/_api_typing.py,sha256=Jj_E61r35EgecWmBvAzpASV4qub5aQI_O4aL-ngEvQ8,23028
+array_api_typing/typing_extra/_at.py,sha256=S7_YjOwR3a8olZWgwpLDFEfnekRufRqrtfiMLwrWjgo,2202
+xbarray/__init__.py,sha256=GLdMGsm5_CnATyh19At20Hmb1MsXZvPxNvg70bHYChk,19
+xbarray/base/__init__.py,sha256=ERmmOxz_9mUkIuccNbzUa5Y6gVLLVDdyc4cCxbCCUbY,20
+xbarray/base/base.py,sha256=hqlizpZOJua9xJfPfc-RR6IAadp-OLPVnBO5cbVVHKQ,5929
+xbarray/common/implementations.py,sha256=FyZaQZcVdSnH6CUA1LVbUJAK-EAlkz2fDfOOuVyku6Y,2655
+xbarray/jax/__init__.py,sha256=J7yN6ytn1fm2qKHZVm9OTt30ItcKMnKnuTxbZgnGKaE,662
+xbarray/jax/_extra.py,sha256=JHZBaGZUXPLCRFdaDBXUjwGD-D4DlLYI0NsxAlGs3FE,2504
+xbarray/jax/_typing.py,sha256=U9BUxHNEjFB0LHF1KMrFLbh6E5kVvpAF8bZUbLNf25E,278
+xbarray/jax/random.py,sha256=owIxqs67MOMXaMN8uwmCu9aLd1Yss24sl9rib4R5Qec,3325
+xbarray/numpy/__init__.py,sha256=trZhNZYIHe0QkQYRB1ZJW-OlteYf_9RVUR3yks6S2iI,523
+xbarray/numpy/_extra.py,sha256=wbzpcVL4UV5TVQlmO73sUwWcBgMiVXFvfCLlCuiB0Uk,1982
+xbarray/numpy/_typing.py,sha256=pgjLLAipwFsIk0QdgrAA4PEvjF_ugHbzfSTbLpA__6o,241
+xbarray/numpy/random.py,sha256=8BaxX_cpzvF7iG0yR0ws3IbG14OJPsX60Gv-1AStplA,2643
+xbarray/pytorch/__init__.py,sha256=wfUkSyrJpDlO2Zggza7HM5SCdYKUBfySO3K71itOHkA,526
+xbarray/pytorch/_extra.py,sha256=VawcjDe0dxiIzvTjIncGxwPPoTw7hWZAPOlZSdDSKl0,2867
+xbarray/pytorch/_typing.py,sha256=qSnNZD3IpgSoOTl3TBRBv2ifSAHOw0aR9uyzfV5KYVw,204
+xbarray/pytorch/random.py,sha256=OJ-YvS2owX5LWdsIzXW-N2YqPhN84Rroco7O5qIMmrg,2670
+xbarray-0.0.1a1.dist-info/licenses/LICENSE,sha256=6P0HCOancSfch0dNycuDIe8_qwS0Id97Ih_8hjJ2PFI,1067
+xbarray-0.0.1a1.dist-info/METADATA,sha256=u2-YoX051dU1Po1mFspLSuIEMmlPdhlWOlNwvbp-8PE,359
+xbarray-0.0.1a1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+xbarray-0.0.1a1.dist-info/top_level.txt,sha256=VriXuFyU48Du4HQMzROSArhwqB6EZYY0n0mipgUqB9A,25
+xbarray-0.0.1a1.dist-info/RECORD,,

xbarray-0.0.1a1.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (80.9.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

xbarray-0.0.1a1.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 Yunhao Cao
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

xbarray-0.0.1a1.dist-info/top_level.txt

@@ -0,0 +1,2 @@
+array_api_typing
+xbarray