Spaces:
Sleeping
Sleeping
| """Utility functions for the Spinor library.""" | |
| from typing import Optional, Sequence, Union | |
| import torch | |
| import numpy as np | |
| TensorLike = Union[torch.Tensor, np.ndarray, list, tuple, float, complex] | |
| # Threshold below which a norm is treated as zero (used to avoid divide-by-zero). | |
| DEFAULT_EPS = 1e-12 | |
| def get_backend(tensor: object) -> str: | |
| """Detect the backend of an array-like object ("torch" or "numpy").""" | |
| if isinstance(tensor, torch.Tensor): | |
| return "torch" | |
| if isinstance(tensor, np.ndarray): | |
| return "numpy" | |
| raise TypeError(f"Unsupported tensor type: {type(tensor)}") | |
| def ensure_torch_tensor( | |
| data: TensorLike, | |
| dtype: Optional[torch.dtype] = None, | |
| device: Optional[Union[str, torch.device]] = None, | |
| ) -> torch.Tensor: | |
| """ | |
| Convert input data to a PyTorch tensor. | |
| Existing tensors are returned as-is (so autograd leaves are preserved), | |
| NumPy arrays are wrapped without copying (memory is shared), and anything | |
| else goes through ``torch.as_tensor``. | |
| """ | |
| if isinstance(data, torch.Tensor): | |
| tensor = data | |
| elif isinstance(data, np.ndarray): | |
| tensor = torch.from_numpy(data) | |
| else: | |
| tensor = torch.as_tensor(data) | |
| if dtype is not None or device is not None: | |
| tensor = tensor.to(dtype=dtype, device=device) | |
| return tensor | |
| def validate_spinor_dimension(tensor: torch.Tensor, dim: int = 2) -> None: | |
| """Validate that the last dimension matches the expected spinor dimension.""" | |
| if tensor.ndim == 0: | |
| raise ValueError("Spinor tensor must have at least one dimension") | |
| if tensor.shape[-1] != dim: | |
| raise ValueError( | |
| f"Last dimension must be {dim} for spinors, got {tensor.shape[-1]}" | |
| ) | |
| def check_broadcastable(a: torch.Tensor, b: torch.Tensor) -> torch.Size: | |
| """Return the broadcast shape of two tensors, raising ValueError if impossible.""" | |
| try: | |
| return torch.broadcast_shapes(a.shape, b.shape) | |
| except RuntimeError as exc: | |
| raise ValueError( | |
| f"Shapes {tuple(a.shape)} and {tuple(b.shape)} are not broadcastable" | |
| ) from exc | |
| def real_dtype_of(dtype: torch.dtype) -> torch.dtype: | |
| """Return the real dtype underlying a (possibly complex) dtype.""" | |
| if dtype == torch.complex64: | |
| return torch.float32 | |
| if dtype == torch.complex128: | |
| return torch.float64 | |
| if dtype == torch.complex32: | |
| return torch.float16 | |
| return dtype | |
| def complex_dtype_of(dtype: torch.dtype) -> torch.dtype: | |
| """Return the complex dtype whose real part has the given precision.""" | |
| if dtype.is_complex: | |
| return dtype | |
| if dtype == torch.float64: | |
| return torch.complex128 | |
| # float32 and everything lower-precision (or non-floating) map to complex64 | |
| return torch.complex64 | |
| def resolve_dtype_and_device( | |
| inputs: Sequence[object], | |
| dtype: Optional[torch.dtype], | |
| device: Optional[Union[str, torch.device]], | |
| ) -> "tuple[torch.dtype, torch.device]": | |
| """ | |
| Pick the real working dtype and device for a computation. | |
| If ``dtype`` is given it wins (a complex dtype is mapped to its real part). | |
| Otherwise, use the highest-precision floating dtype among any tensor inputs, | |
| falling back to the torch default. The device comes from ``device`` or the | |
| first tensor input, falling back to CPU. | |
| """ | |
| if dtype is not None: | |
| real_dtype = real_dtype_of(dtype) | |
| else: | |
| real_dtype = None | |
| for x in inputs: | |
| if isinstance(x, torch.Tensor) and x.is_floating_point(): | |
| real_dtype = x.dtype if real_dtype is None else torch.promote_types(real_dtype, x.dtype) | |
| if real_dtype is None: | |
| real_dtype = torch.get_default_dtype() | |
| if device is None: | |
| for x in inputs: | |
| if isinstance(x, torch.Tensor): | |
| device = x.device | |
| break | |
| return real_dtype, torch.device(device) if device is not None else torch.device("cpu") | |