from abc import ABC, abstractmethod
from typing import Any, Tuple, List, Optional, Union, Sequence, ContextManager, Dict
[docs]
class TensorOps(ABC):
"""
Abstract base class for tensor operations across different backends.
"""
[docs]
@abstractmethod
def create_tensor(self, data: Any, **kwargs) -> Any:
pass
[docs]
@abstractmethod
def shape(self, tensor: Any) -> Any:
pass
[docs]
@abstractmethod
def from_numpy(self, array: Any) -> Any:
pass
[docs]
@abstractmethod
def to_numpy(self, tensor: Any) -> Any:
pass
[docs]
@abstractmethod
def no_grad(self) -> ContextManager:
pass
[docs]
@abstractmethod
def zeros(self, shape: Tuple[int, ...], **kwargs) -> Any:
pass
[docs]
@abstractmethod
def ones(self, shape: Tuple[int, ...], **kwargs) -> Any:
pass
[docs]
@abstractmethod
def eye(self, n: int, **kwargs) -> Any:
pass
[docs]
@abstractmethod
def arange(self, start: int, end: int, step: int = 1, **kwargs) -> Any:
pass
[docs]
@abstractmethod
def linspace(self, start: float, end: float, num: int, **kwargs) -> Any:
pass
[docs]
@abstractmethod
def zeros_like(self, tensor: Any, **kwargs) -> Any:
pass
[docs]
@abstractmethod
def ones_like(self, tensor: Any, **kwargs) -> Any:
pass
[docs]
@abstractmethod
def sqrt(self, tensor: Any) -> Any:
pass
[docs]
@abstractmethod
def stack(self, tensors: List[Any], dim: int = 0) -> Any:
pass
[docs]
@abstractmethod
def cat(self, tensors: List[Any], dim: int = 0) -> Any:
pass
[docs]
@abstractmethod
def reshape(self, tensor: Any, shape: Tuple[int, ...]) -> Any:
pass
[docs]
@abstractmethod
def repeat(self, tensor: Any, repeats: Union[int, Tuple[int, ...]], dim: int = 0) -> Any:
pass
[docs]
@abstractmethod
def tile(self, tensor: Any, reps: Union[int, Tuple[int, ...]]) -> Any:
pass
[docs]
@abstractmethod
def clip(self, tensor: Any, min_val: float, max_val: float) -> Any:
pass
[docs]
@abstractmethod
def unsqueeze(self, tensor: Any, dim: int) -> Any:
pass
[docs]
@abstractmethod
def expand(self, tensor: Any, *sizes: int) -> Any:
pass
[docs]
@abstractmethod
def gather(self, tensor: Any, dim: int, index: Any) -> Any:
pass
[docs]
@abstractmethod
def squeeze(self, tensor: Any, dim: Optional[int] = None) -> Any:
pass
[docs]
@abstractmethod
def transpose(self, tensor: Any, *args, **kwargs) -> Any:
pass
[docs]
@abstractmethod
def matmul(self, a: Any, b: Any) -> Any:
pass
[docs]
@abstractmethod
def inverse(self, tensor: Any) -> Any:
pass
[docs]
@abstractmethod
def mean(
self,
tensor: Any,
dim: Optional[int] = None,
keepdims: bool = False,
axis: Optional[int] = None,
) -> Any:
pass
[docs]
@abstractmethod
def sum(
self,
tensor: Any,
dim: Optional[int] = None,
keepdims: bool = False,
axis: Optional[int] = None,
) -> Any:
pass
[docs]
@abstractmethod
def max(
self,
tensor: Any,
dim: Optional[int] = None,
keepdims: bool = False,
axis: Optional[int] = None,
) -> Any:
pass
[docs]
@abstractmethod
def min(
self,
tensor: Any,
dim: Optional[int] = None,
keepdims: bool = False,
axis: Optional[int] = None,
) -> Any:
pass
[docs]
@abstractmethod
def minimum(self, a: Any, b: Any) -> Any:
pass
[docs]
@abstractmethod
def maximum(self, a: Any, b: Any) -> Any:
pass
[docs]
@abstractmethod
def where(self, condition: Any, x: Any, y: Any) -> Any:
pass
[docs]
@abstractmethod
def norm(
self,
tensor: Any,
dim: Optional[int] = None,
axis: Optional[int] = None,
keepdims: bool = False,
ord: Union[int, float, str] = 2,
) -> Any:
pass
[docs]
def tensor(self, data: Any, **kwargs) -> Any:
"""Alias for create_tensor"""
return self.create_tensor(data, **kwargs)
# --- Arithmetic ---
[docs]
@abstractmethod
def add(self, a: Any, b: Any) -> Any: pass
[docs]
@abstractmethod
def subtract(self, a: Any, b: Any) -> Any: pass
[docs]
@abstractmethod
def multiply(self, a: Any, b: Any) -> Any: pass
[docs]
@abstractmethod
def divide(self, a: Any, b: Any) -> Any: pass
[docs]
@abstractmethod
def power(self, a: Any, b: Any) -> Any: pass
# --- Math ---
[docs]
@abstractmethod
def sin(self, tensor: Any) -> Any: pass
[docs]
@abstractmethod
def cos(self, tensor: Any) -> Any: pass
[docs]
@abstractmethod
def exp(self, tensor: Any) -> Any: pass
[docs]
@abstractmethod
def log(self, tensor: Any) -> Any: pass
[docs]
@abstractmethod
def abs(self, tensor: Any) -> Any: pass
[docs]
@abstractmethod
def sign(self, tensor: Any) -> Any: pass
# --- Activation ---
[docs]
@abstractmethod
def relu(self, tensor: Any) -> Any: pass
[docs]
@abstractmethod
def sigmoid(self, tensor: Any) -> Any: pass
[docs]
@abstractmethod
def tanh(self, tensor: Any) -> Any: pass
[docs]
@abstractmethod
def softmax(self, tensor: Any, dim: int = -1) -> Any: pass
# --- NN Operations ---
[docs]
@abstractmethod
def dropout(self, tensor: Any, p: float = 0.5, training: bool = True, **kwargs) -> Any: pass
[docs]
@abstractmethod
def batch_norm(self, tensor: Any, mean: Any = None, var: Any = None, weight: Any = None, bias: Any = None, training: bool = False, momentum: float = 0.1, eps: float = 1e-5) -> Any: pass
[docs]
@abstractmethod
def layer_norm(self, tensor: Any, normalized_shape: Any = None, weight: Any = None, bias: Any = None, eps: float = 1e-5) -> Any: pass
[docs]
@abstractmethod
def convolve(self, input: Any, weight: Any, bias: Optional[Any] = None, stride: int = 1, padding: int = 0, dilation: int = 1, groups: int = 1) -> Any: pass
[docs]
@abstractmethod
def max_pool(self, input: Any, kernel_size: int, stride: Optional[int] = None, padding: int = 0) -> Any: pass
[docs]
@abstractmethod
def avg_pool(self, input: Any, kernel_size: int, stride: Optional[int] = None, padding: int = 0) -> Any: pass
# --- Gradient/Autograd ---
[docs]
@abstractmethod
def backward(self, tensor: Any, grad_tensor: Optional[Any] = None) -> Any: pass
[docs]
@abstractmethod
def grad(self, tensor: Any) -> Any: pass
# --- Random ---
[docs]
@abstractmethod
def random_normal(self, shape: Tuple[int, ...], mean: float = 0.0, std: float = 1.0, **kwargs) -> Any: pass
# --- Loss ---
[docs]
@abstractmethod
def mse_loss(self, input: Any, target: Any, reduction: str = 'mean') -> Any: pass
[docs]
@abstractmethod
def cross_entropy_loss(self, input: Any, target: Any, reduction: str = 'mean', **kwargs) -> Any: pass
# --- Device/Misc ---
[docs]
@abstractmethod
def to_device(self, tensor: Any, device: str) -> Any: pass
[docs]
@abstractmethod
def device(self, tensor: Any) -> str: pass
# --- Logic ---
[docs]
@abstractmethod
def equal(self, a: Any, b: Any) -> Any: pass
[docs]
@abstractmethod
def greater(self, a: Any, b: Any) -> Any: pass
[docs]
@abstractmethod
def less(self, a: Any, b: Any) -> Any: pass
[docs]
@abstractmethod
def logical_and(self, a: Any, b: Any) -> Any: pass
[docs]
@abstractmethod
def logical_or(self, a: Any, b: Any) -> Any: pass
[docs]
@abstractmethod
def logical_not(self, a: Any) -> Any: pass
# --- FFT ---
[docs]
@abstractmethod
def fft(self, tensor: Any, n: Optional[int] = None, dim: int = -1, norm: Optional[str] = None) -> Any: pass
[docs]
@abstractmethod
def ifft(self, tensor: Any, n: Optional[int] = None, dim: int = -1, norm: Optional[str] = None) -> Any:
pass
# --- Extra Utils for Tests ---
[docs]
def concatenate(self, tensors: List[Any], dim: int = 0) -> Any:
return self.cat(tensors, dim)
[docs]
@abstractmethod
def index(self, tensor: Any, index: Any) -> Any:
pass
[docs]
@abstractmethod
def slice(self, tensor: Any, start: int, end: int, dim: int = 0) -> Any:
pass
[docs]
@abstractmethod
def sgd_step(self, tensor: Any, lr: float) -> Any:
pass
[docs]
@abstractmethod
def adam_step(self, tensor: Any, lr: float) -> Any:
pass
[docs]
def switch_backend(self, backend: Any) -> bool:
# Default implementation calling global manager
from hpfracc.ml.backends import get_backend_manager
return get_backend_manager().switch_backend(backend)
[docs]
@abstractmethod
def get_backend_info(self) -> Dict[str, Any]:
pass
[docs]
def enable_profiling(self, enable: bool):
pass
[docs]
def get_profile_results(self) -> Dict[str, Any]:
return {}
[docs]
def clear_cache(self):
pass
[docs]
def get_memory_usage(self) -> Dict[str, Any]:
return {}