PyNorch/norch/tensor.py

1080 lines
39 KiB
Python
Raw Permalink Normal View History

2024-04-26 18:38:07 -03:00
import ctypes
2024-04-30 00:31:46 -03:00
import os
2024-04-30 20:09:08 -03:00
from .autograd.functions import *
2024-04-26 18:38:07 -03:00
class CTensor(ctypes.Structure):
_fields_ = [
('data', ctypes.POINTER(ctypes.c_float)),
('strides', ctypes.POINTER(ctypes.c_int)),
('shape', ctypes.POINTER(ctypes.c_int)),
('ndim', ctypes.c_int),
('size', ctypes.c_int),
2024-04-29 02:20:31 -03:00
('device', ctypes.c_char_p)
2024-04-26 18:38:07 -03:00
]
class Tensor:
2024-05-11 01:45:59 -03:00
module_dir = os.path.dirname(os.path.abspath(__file__))
_C = ctypes.CDLL(os.path.join(module_dir, "libtensor.so"))
2024-04-26 18:38:07 -03:00
2024-04-30 20:09:08 -03:00
def __init__(self, data=None, device="cpu", requires_grad=False):
2024-04-27 17:24:55 -03:00
if data != None:
2024-05-18 17:59:14 -03:00
if isinstance(data, (float, int)):
data = [data]
2024-04-27 17:24:55 -03:00
data, shape = self.flatten(data)
2024-05-18 17:59:14 -03:00
self.shape = shape.copy()
2024-06-05 20:24:09 -03:00
self._data_ctype = (ctypes.c_float * len(data))(*data.copy())
self._shape_ctype = (ctypes.c_int * len(shape))(*shape.copy())
self._ndim_ctype = ctypes.c_int(len(shape))
self._device_ctype = device.encode('utf-8')
2024-04-29 02:20:31 -03:00
2024-04-27 17:24:55 -03:00
self.ndim = len(shape)
2024-04-29 02:20:31 -03:00
self.device = device
2024-04-27 17:24:55 -03:00
self.numel = 1
for s in self.shape:
self.numel *= s
2024-04-30 20:09:08 -03:00
self.requires_grad = requires_grad
self.hooks = []
2024-04-30 20:09:08 -03:00
self.grad = None
self.grad_fn = None
2024-04-29 02:20:31 -03:00
Tensor._C.create_tensor.argtypes = [ctypes.POINTER(ctypes.c_float), ctypes.POINTER(ctypes.c_int), ctypes.c_int, ctypes.c_char_p]
2024-04-27 17:24:55 -03:00
Tensor._C.create_tensor.restype = ctypes.POINTER(CTensor)
2024-04-29 02:20:31 -03:00
2024-04-27 17:24:55 -03:00
self.tensor = Tensor._C.create_tensor(
2024-06-05 20:24:09 -03:00
self._data_ctype,
self._shape_ctype,
self._ndim_ctype,
self._device_ctype
2024-04-27 17:24:55 -03:00
)
else:
self.tensor = None,
self.shape = None,
2024-04-29 02:20:31 -03:00
self.ndim = None,
self.device = device
2024-04-30 20:09:08 -03:00
self.requires_grad = None
self.hooks = []
2024-04-30 20:09:08 -03:00
self.grad = None
self.grad_fn = None
2024-04-26 18:38:07 -03:00
def flatten(self, nested_list):
2024-04-28 18:43:07 -03:00
def flatten_recursively(nested_list):
flat_data = []
shape = []
if isinstance(nested_list, list):
for sublist in nested_list:
inner_data, inner_shape = flatten_recursively(sublist)
flat_data.extend(inner_data)
shape.append(len(nested_list))
shape.extend(inner_shape)
else:
flat_data.append(nested_list)
return flat_data, shape
flat_data, shape = flatten_recursively(nested_list)
2024-04-26 18:38:07 -03:00
return flat_data, shape
2024-06-05 14:31:48 -03:00
def __del__(self):
2024-06-05 20:24:09 -03:00
if hasattr(self, '_data_ctype') and self._data_ctype is not None:
2024-06-05 21:04:09 -03:00
2024-06-05 20:24:09 -03:00
Tensor._C.delete_strides.argtypes = [ctypes.POINTER(CTensor)]
Tensor._C.delete_strides.restype = None
Tensor._C.delete_strides(self.tensor)
Tensor._C.delete_device.argtypes = [ctypes.POINTER(CTensor)]
Tensor._C.delete_device.restype = None
Tensor._C.delete_device(self.tensor)
2024-06-05 21:04:09 -03:00
Tensor._C.delete_tensor.argtypes = [ctypes.POINTER(CTensor)]
Tensor._C.delete_tensor.restype = None
Tensor._C.delete_tensor(self.tensor)
2024-06-05 20:24:09 -03:00
elif self.tensor is not None:
2024-06-05 21:04:09 -03:00
Tensor._C.delete_strides.argtypes = [ctypes.POINTER(CTensor)]
Tensor._C.delete_strides.restype = None
Tensor._C.delete_strides(self.tensor)
Tensor._C.delete_data.argtypes = [ctypes.POINTER(CTensor)]
Tensor._C.delete_data.restype = None
Tensor._C.delete_data(self.tensor)
Tensor._C.delete_shape.argtypes = [ctypes.POINTER(CTensor)]
Tensor._C.delete_shape.restype = None
Tensor._C.delete_shape(self.tensor)
Tensor._C.delete_device.argtypes = [ctypes.POINTER(CTensor)]
Tensor._C.delete_device.restype = None
Tensor._C.delete_device(self.tensor)
2024-06-05 20:24:09 -03:00
2024-06-05 14:31:48 -03:00
Tensor._C.delete_tensor.argtypes = [ctypes.POINTER(CTensor)]
Tensor._C.delete_tensor.restype = None
Tensor._C.delete_tensor(self.tensor)
def __setattr__(self, name, value):
if name == 'grad':
for hook in self.hooks:
value = hook(value)
super().__setattr__(name, value)
def register_hook(self, function):
self.hooks.append(function)
2024-04-30 20:09:08 -03:00
def ones_like(self):
Tensor._C.ones_like_tensor.argtypes = [ctypes.POINTER(CTensor)]
Tensor._C.ones_like_tensor.restype = ctypes.POINTER(CTensor)
Tensor._C.ones_like_tensor(self.tensor)
result_tensor_ptr = Tensor._C.ones_like_tensor(self.tensor)
result_data = Tensor()
result_data.tensor = result_tensor_ptr
result_data.shape = self.shape.copy()
result_data.ndim = self.ndim
result_data.device = self.device
result_data.numel = self.numel
2024-04-30 20:09:08 -03:00
return result_data
def zeros_like(self):
Tensor._C.zeros_like_tensor.argtypes = [ctypes.POINTER(CTensor)]
Tensor._C.zeros_like_tensor.restype = ctypes.POINTER(CTensor)
Tensor._C.zeros_like_tensor(self.tensor)
2024-05-01 12:11:50 -03:00
result_tensor_ptr = Tensor._C.zeros_like_tensor(self.tensor)
2024-04-30 20:09:08 -03:00
result_data = Tensor()
result_data.tensor = result_tensor_ptr
result_data.shape = self.shape.copy()
result_data.ndim = self.ndim
result_data.device = self.device
result_data.numel = self.numel
2024-04-30 20:09:08 -03:00
return result_data
2024-05-07 12:48:06 -03:00
def reshape(self, new_shape):
2024-05-18 17:59:14 -03:00
# Calculate the total number of elements in the tensor
total_elements = self.numel
# Check for the presence of -1 in new_shape
if new_shape.count(-1) > 1:
raise ValueError("Only one dimension can be inferred (set to -1).")
inferred_dim = None
known_dims_product = 1
for dim in new_shape:
if dim == -1:
inferred_dim = dim
else:
known_dims_product *= dim
# Calculate the inferred dimension if -1 is present
if inferred_dim == -1:
inferred_dim_size = total_elements // known_dims_product
new_shape = [inferred_dim_size if dim == -1 else dim for dim in new_shape]
2024-04-28 18:43:07 -03:00
new_shape_ctype = (ctypes.c_int * len(new_shape))(*new_shape)
new_ndim_ctype = ctypes.c_int(len(new_shape))
Tensor._C.reshape_tensor.argtypes = [ctypes.POINTER(CTensor), ctypes.POINTER(ctypes.c_int), ctypes.c_int]
2024-05-01 12:11:50 -03:00
Tensor._C.reshape_tensor.restype = ctypes.POINTER(CTensor)
result_tensor_ptr = Tensor._C.reshape_tensor(self.tensor, new_shape_ctype, new_ndim_ctype)
result_data = Tensor()
result_data.tensor = result_tensor_ptr
result_data.shape = new_shape.copy()
result_data.ndim = len(new_shape)
result_data.device = self.device
result_data.numel = self.numel
2024-04-28 18:43:07 -03:00
2024-05-01 12:11:50 -03:00
result_data.requires_grad = self.requires_grad
2024-05-07 12:48:06 -03:00
if result_data.requires_grad:
result_data.grad_fn = ReshapeBackward(self)
2024-04-28 18:43:07 -03:00
2024-05-03 13:18:37 -03:00
return result_data
2024-05-21 14:21:35 -03:00
def unsqueeze(self, dim):
2024-05-21 14:32:14 -03:00
if dim < 0:
dim = self.ndim + dim + 1
2024-05-21 14:21:35 -03:00
# Ensure the dimension is valid
2024-05-21 14:32:14 -03:00
if dim > self.ndim:
2024-05-21 14:21:35 -03:00
raise ValueError("Dimension out of range (expected to be in range of [0, {0}], but got {1})".format(self.ndim, dim))
# Create the new shape with an extra dimension of size 1
new_shape = self.shape[:dim] + [1] + self.shape[dim:]
return self.reshape(new_shape)
2024-05-21 14:43:16 -03:00
def squeeze(self, dim=None):
if dim is not None:
if dim < 0:
dim = self.ndim + dim
# Ensure the dimension is valid
if dim >= self.ndim or dim < 0:
raise ValueError("Dimension out of range (expected to be in range of [0, {0}), but got {1})".format(self.ndim, dim))
# Only squeeze the specified dimension if its size is 1
if self.shape[dim] != 1:
2024-05-22 10:34:16 -03:00
return self
#raise ValueError("Dimension {0} does not have size 1 and cannot be squeezed".format(dim))
2024-05-21 14:43:16 -03:00
# Create the new shape without the specified dimension
new_shape = self.shape[:dim] + self.shape[dim+1:]
else:
# Create the new shape by removing all dimensions of size 1
new_shape = [s for s in self.shape if s != 1]
return self.reshape(new_shape)
2024-04-29 02:20:31 -03:00
def to(self, device):
2024-05-23 19:15:30 -03:00
device = str(device)
2024-04-29 02:20:31 -03:00
self.device = device
2024-04-29 19:21:10 -03:00
self.device_ctype = self.device.encode('utf-8')
2024-04-29 02:20:31 -03:00
Tensor._C.to_device.argtypes = [ctypes.POINTER(CTensor), ctypes.c_char_p]
Tensor._C.to_device.restype = None
Tensor._C.to_device(self.tensor, self.device_ctype)
2024-04-29 10:12:01 -03:00
return self
2024-04-30 20:09:08 -03:00
def backward(self, gradient=None):
if not self.requires_grad:
return
2024-04-30 20:09:08 -03:00
if gradient is None:
2024-05-01 00:20:11 -03:00
if self.shape == [1]:
2024-05-23 12:07:17 -03:00
gradient = Tensor([1]).to(self.device)
2024-05-01 00:20:11 -03:00
else:
2024-04-30 20:34:58 -03:00
raise RuntimeError("Gradient argument must be specified for non-scalar tensors.")
2024-05-01 00:20:11 -03:00
2024-05-07 01:00:08 -03:00
stack = [(self, gradient)]
visited = set()
while stack:
tensor, grad = stack.pop()
if tensor.grad is None:
tensor.grad = grad
else:
tensor.grad += grad
2024-04-30 20:09:08 -03:00
# Propagate gradients to inputs if not a leaf tensor
if tensor.grad_fn is not None:
grads = tensor.grad_fn.backward(grad)
for tensor, grad in zip(tensor.grad_fn.input, grads):
if isinstance(tensor, Tensor) and tensor not in visited:
stack.append((tensor, grad))
visited.add(tensor)
2024-04-30 20:09:08 -03:00
def zero_grad(self):
2024-05-07 00:10:25 -03:00
self.grad = None
2024-04-29 10:12:01 -03:00
2024-05-07 01:19:12 -03:00
def detach(self):
self.grad = None
self.grad_fn = None
2024-04-26 18:38:07 -03:00
def __getitem__(self, indices):
2024-05-10 02:01:09 -03:00
if isinstance(indices, int):
indices = [indices]
2024-04-26 18:38:07 -03:00
if len(indices) != self.ndim:
raise ValueError("Number of indices must match the number of dimensions")
Tensor._C.get_item.argtypes = [ctypes.POINTER(CTensor), ctypes.POINTER(ctypes.c_int)]
Tensor._C.get_item.restype = ctypes.c_float
indices = (ctypes.c_int * len(indices))(*indices)
value = Tensor._C.get_item(self.tensor, indices)
return value
2024-04-28 18:43:07 -03:00
2024-04-26 18:38:07 -03:00
def __str__(self):
2024-04-28 18:43:07 -03:00
def print_recursively(tensor, depth, index):
if depth == tensor.ndim - 1:
result = ""
for i in range(tensor.shape[-1]):
index[-1] = i
result += str(tensor[tuple(index)]) + ", "
return result.strip()
else:
result = ""
if depth > 0:
result += "\n" + " " * ((depth - 1) * 4)
for i in range(tensor.shape[depth]):
index[depth] = i
result += "["
result += print_recursively(tensor, depth + 1, index) + "],"
if i < tensor.shape[depth] - 1:
result += "\n" + " " * (depth * 4)
return result.strip(",")
index = [0] * self.ndim
result = "tensor(["
result += print_recursively(self, 0, index)
2024-04-30 20:09:08 -03:00
result += f"""], device="{self.device}", requires_grad={self.requires_grad})"""
2024-04-28 18:43:07 -03:00
return result
2024-04-26 18:38:07 -03:00
def __repr__(self):
return self.__str__()
def __add__(self, other):
2024-05-06 10:18:12 -03:00
if isinstance(other, (int, float)):
other = other * self.ones_like()
broadcasted_shape_add = []
2024-05-16 16:35:55 -03:00
# Function to determine if broadcasting is needed and get the broadcasted shape
def broadcast_shape(shape1, shape2):
if shape1 == shape2:
return shape1, False
max_len = max(len(shape1), len(shape2))
shape1 = [1] * (max_len - len(shape1)) + shape1
shape2 = [1] * (max_len - len(shape2)) + shape2
2024-05-16 16:35:55 -03:00
for dim1, dim2 in zip(shape1, shape2):
if dim1 != dim2 and dim1 != 1 and dim2 != 1:
raise ValueError("Shapes are not compatible for broadcasting")
broadcasted_shape_add.append(max(dim1, dim2))
return broadcasted_shape_add, True
2024-05-16 16:35:55 -03:00
broadcasted_shape_add, needs_broadcasting = broadcast_shape(self.shape, other.shape)
2024-05-16 16:35:55 -03:00
if needs_broadcasting:
# Call add_broadcasted_tensor if broadcasting is needed
if other.ndim == self.ndim - 1:
other = other.reshape([1] + other.shape)
elif self.ndim == other.ndim - 1:
self = self.reshape([1] + self.shape)
2024-05-16 16:35:55 -03:00
Tensor._C.add_broadcasted_tensor.argtypes = [ctypes.POINTER(CTensor), ctypes.POINTER(CTensor)]
Tensor._C.add_broadcasted_tensor.restype = ctypes.POINTER(CTensor)
result_tensor_ptr = Tensor._C.add_broadcasted_tensor(self.tensor, other.tensor)
result_data = Tensor()
result_data.tensor = result_tensor_ptr
result_data.shape = broadcasted_shape_add.copy()
result_data.ndim = len(broadcasted_shape_add)
2024-05-16 16:35:55 -03:00
result_data.device = self.device
result_data.numel = 1
for s in result_data.shape:
result_data.numel *= s
2024-05-16 16:35:55 -03:00
result_data.requires_grad = self.requires_grad or other.requires_grad
if result_data.requires_grad:
result_data.grad_fn = AddBroadcastedBackward(self, other)
2024-05-06 10:18:12 -03:00
2024-05-16 16:35:55 -03:00
else:
# Call add_tensor if shapes are identical
Tensor._C.add_tensor.argtypes = [ctypes.POINTER(CTensor), ctypes.POINTER(CTensor)]
Tensor._C.add_tensor.restype = ctypes.POINTER(CTensor)
2024-04-26 18:38:07 -03:00
2024-05-16 16:35:55 -03:00
result_tensor_ptr = Tensor._C.add_tensor(self.tensor, other.tensor)
2024-04-26 18:38:07 -03:00
2024-05-16 16:35:55 -03:00
result_data = Tensor()
result_data.tensor = result_tensor_ptr
result_data.shape = self.shape.copy()
result_data.ndim = self.ndim
2024-05-16 16:35:55 -03:00
result_data.device = self.device
result_data.numel = self.numel # Update this to calculate the correct number of elements if broadcasting
result_data.requires_grad = self.requires_grad or other.requires_grad
if result_data.requires_grad:
result_data.grad_fn = AddBackward(self, other)
2024-04-26 18:38:07 -03:00
2024-04-27 17:24:55 -03:00
return result_data
2024-05-16 16:35:55 -03:00
2024-04-27 17:24:55 -03:00
2024-05-06 10:24:43 -03:00
def __radd__(self, other):
if isinstance(other, (int, float)):
other = other * self.ones_like()
if self.shape != other.shape:
raise ValueError("Tensors must have the same shape for addition")
Tensor._C.add_tensor.argtypes = [ctypes.POINTER(CTensor), ctypes.POINTER(CTensor)]
Tensor._C.add_tensor.restype = ctypes.POINTER(CTensor)
result_tensor_ptr = Tensor._C.add_tensor(other.tensor, self.tensor)
result_data = Tensor()
result_data.tensor = result_tensor_ptr
result_data.shape = self.shape.copy()
result_data.ndim = self.ndim
result_data.device = self.device
result_data.numel = self.numel
2024-05-06 10:24:43 -03:00
result_data.requires_grad = self.requires_grad or other.requires_grad
if result_data.requires_grad:
result_data.grad_fn = AddBackward(other, self)
return result_data
2024-04-27 17:24:55 -03:00
def __sub__(self, other):
2024-05-06 10:18:12 -03:00
if isinstance(other, (int, float)):
other = other * self.ones_like()
2024-05-06 10:24:43 -03:00
broadcasted_shape_sub = []
2024-05-16 17:03:09 -03:00
# Function to determine if broadcasting is needed and get the broadcasted shape
def broadcast_shape(shape1, shape2):
if shape1 == shape2:
return shape1, False
max_len = max(len(shape1), len(shape2))
shape1 = [1] * (max_len - len(shape1)) + shape1
shape2 = [1] * (max_len - len(shape2)) + shape2
2024-05-16 17:03:09 -03:00
for dim1, dim2 in zip(shape1, shape2):
if dim1 != dim2 and dim1 != 1 and dim2 != 1:
raise ValueError("Shapes are not compatible for broadcasting")
broadcasted_shape_sub.append(max(dim1, dim2))
return broadcasted_shape_sub, True
2024-05-16 17:03:09 -03:00
broadcasted_shape_sub, needs_broadcasting = broadcast_shape(self.shape, other.shape)
2024-05-16 17:03:09 -03:00
if needs_broadcasting:
if other.ndim == self.ndim - 1:
other = other.reshape([1] + other.shape)
elif self.ndim == other.ndim - 1:
self = self.reshape([1] + self.shape)
2024-05-16 17:03:09 -03:00
# Call add_broadcasted_tensor if broadcasting is needed
Tensor._C.sub_broadcasted_tensor.argtypes = [ctypes.POINTER(CTensor), ctypes.POINTER(CTensor)]
Tensor._C.sub_broadcasted_tensor.restype = ctypes.POINTER(CTensor)
result_tensor_ptr = Tensor._C.sub_broadcasted_tensor(self.tensor, other.tensor)
result_data = Tensor()
result_data.tensor = result_tensor_ptr
result_data.shape = broadcasted_shape_sub.copy()
result_data.ndim = len(broadcasted_shape_sub)
2024-05-16 17:03:09 -03:00
result_data.device = self.device
result_data.numel = self.numel # Update this to calculate the correct number of elements if broadcasting
result_data.requires_grad = self.requires_grad or other.requires_grad
if result_data.requires_grad:
result_data.grad_fn = SubBroadcastedBackward(self, other)
2024-04-27 17:24:55 -03:00
2024-05-16 17:03:09 -03:00
else:
# Call add_tensor if shapes are identical
Tensor._C.sub_tensor.argtypes = [ctypes.POINTER(CTensor), ctypes.POINTER(CTensor)]
Tensor._C.sub_tensor.restype = ctypes.POINTER(CTensor)
2024-04-27 17:24:55 -03:00
2024-05-16 17:03:09 -03:00
result_tensor_ptr = Tensor._C.sub_tensor(self.tensor, other.tensor)
2024-04-27 17:24:55 -03:00
2024-05-16 17:03:09 -03:00
result_data = Tensor()
result_data.tensor = result_tensor_ptr
result_data.shape = self.shape.copy()
result_data.ndim = self.ndim
2024-04-26 18:38:07 -03:00
2024-05-16 17:03:09 -03:00
result_data.device = self.device
result_data.numel = self.numel # Update this to calculate the correct number of elements if broadcasting
result_data.requires_grad = self.requires_grad or other.requires_grad
if result_data.requires_grad:
result_data.grad_fn = SubBackward(self, other)
2024-04-30 20:34:58 -03:00
2024-04-27 17:24:55 -03:00
return result_data
2024-04-28 18:43:07 -03:00
2024-05-06 10:24:43 -03:00
def __rsub__(self, other):
if isinstance(other, (int, float)):
other = other * self.ones_like()
if self.shape != other.shape:
raise ValueError("Tensors must have the same shape for subtraction")
Tensor._C.sub_tensor.argtypes = [ctypes.POINTER(CTensor), ctypes.POINTER(CTensor)]
Tensor._C.sub_tensor.restype = ctypes.POINTER(CTensor)
result_tensor_ptr = Tensor._C.sub_tensor(other.tensor, self.tensor)
result_data = Tensor()
result_data.tensor = result_tensor_ptr
result_data.shape = self.shape.copy()
result_data.ndim = self.ndim
result_data.device = self.device
result_data.numel = self.numel
2024-05-06 10:24:43 -03:00
result_data.requires_grad = self.requires_grad or other.requires_grad
if result_data.requires_grad:
result_data.grad_fn = SubBackward(other, self)
return result_data
2024-04-28 18:43:07 -03:00
def __mul__(self, other):
2024-04-30 13:22:27 -03:00
if isinstance(other, (int, float)):
result_data = Tensor()
result_data.shape = self.shape.copy()
result_data.ndim = self.ndim
result_data.device = self.device
result_data.numel = self.numel
2024-04-30 13:22:27 -03:00
Tensor._C.scalar_mul_tensor.argtypes = [ctypes.POINTER(CTensor), ctypes.c_float]
Tensor._C.scalar_mul_tensor.restype = ctypes.POINTER(CTensor)
2024-04-28 18:43:07 -03:00
2024-04-30 13:22:27 -03:00
result_data.tensor = Tensor._C.scalar_mul_tensor(self.tensor, ctypes.c_float(other))
2024-04-28 18:43:07 -03:00
2024-04-30 20:34:58 -03:00
result_data.requires_grad = self.requires_grad
if result_data.requires_grad:
result_data.grad_fn = ScalarMulBackward(self, other)
2024-04-30 13:22:27 -03:00
return result_data
elif isinstance(other, Tensor):
if self.shape != other.shape:
raise ValueError("Tensors must have the same shape for element-wise multiplication")
2024-04-28 18:43:07 -03:00
2024-04-30 13:22:27 -03:00
Tensor._C.elementwise_mul_tensor.argtypes = [ctypes.POINTER(CTensor), ctypes.POINTER(CTensor)]
Tensor._C.elementwise_mul_tensor.restype = ctypes.POINTER(CTensor)
result_tensor_ptr = Tensor._C.elementwise_mul_tensor(self.tensor, other.tensor)
result_data = Tensor()
result_data.tensor = result_tensor_ptr
result_data.shape = self.shape.copy()
result_data.ndim = self.ndim
result_data.device = self.device
result_data.numel = self.numel
2024-04-30 13:22:27 -03:00
2024-04-30 20:34:58 -03:00
result_data.requires_grad = self.requires_grad or other.requires_grad
if result_data.requires_grad:
result_data.grad_fn = ElementwiseMulBackward(self, other)
2024-04-30 13:22:27 -03:00
return result_data
else:
raise TypeError("Unsupported operand type(s) for *: '{}' and '{}'".format(type(self), type(other)))
def __rmul__(self, other):
return self.__mul__(other)
def __neg__(self):
return self.__mul__(-1)
def __pos__(self):
return self
2024-04-28 18:43:07 -03:00
def __matmul__(self, other):
2024-05-02 20:21:19 -03:00
if self.ndim < 3 and other.ndim == 3:
#broadcasted 2D x 3D matmul
Tensor._C.broadcasted_batched_matmul_tensor.argtypes = [ctypes.POINTER(CTensor), ctypes.POINTER(CTensor)]
Tensor._C.broadcasted_batched_matmul_tensor.restype = ctypes.POINTER(CTensor)
result_tensor_ptr = Tensor._C.broadcasted_batched_matmul_tensor(self.tensor, other.tensor)
result_data = Tensor()
result_data.tensor = result_tensor_ptr
result_data.shape = [other.shape[0], self.shape[0], other.shape[2]]
result_data.ndim = 3
result_data.device = self.device
result_data.numel = 1
for s in result_data.shape:
result_data.numel *= s
2024-05-02 20:21:19 -03:00
elif self.ndim == 3 and other.ndim == 3:
#broadcasted 3D x 3D matmul
2024-04-28 18:43:07 -03:00
2024-05-02 19:13:05 -03:00
Tensor._C.batched_matmul_tensor.argtypes = [ctypes.POINTER(CTensor), ctypes.POINTER(CTensor)]
Tensor._C.batched_matmul_tensor.restype = ctypes.POINTER(CTensor)
result_tensor_ptr = Tensor._C.batched_matmul_tensor(self.tensor, other.tensor)
2024-04-28 18:43:07 -03:00
2024-05-02 19:13:05 -03:00
result_data = Tensor()
result_data.tensor = result_tensor_ptr
2024-05-02 20:21:19 -03:00
result_data.shape = [other.shape[0], self.shape[1], other.shape[2]]
2024-05-02 19:13:05 -03:00
result_data.ndim = 3
result_data.device = self.device
result_data.numel = 1
for s in result_data.shape:
result_data.numel *= s
2024-05-02 19:13:05 -03:00
else:
#2D matmul
if self.ndim != 2 or other.ndim != 2:
raise ValueError("Matrix multiplication requires 2D tensors")
2024-04-28 18:43:07 -03:00
2024-05-02 19:13:05 -03:00
if self.shape[1] != other.shape[0]:
raise ValueError("Incompatible shapes for matrix multiplication")
2024-04-28 18:43:07 -03:00
2024-05-02 19:13:05 -03:00
Tensor._C.matmul_tensor.argtypes = [ctypes.POINTER(CTensor), ctypes.POINTER(CTensor)]
Tensor._C.matmul_tensor.restype = ctypes.POINTER(CTensor)
result_tensor_ptr = Tensor._C.matmul_tensor(self.tensor, other.tensor)
result_data = Tensor()
result_data.tensor = result_tensor_ptr
result_data.shape = [self.shape[0], other.shape[1]]
result_data.ndim = 2
result_data.device = self.device
result_data.numel = 1
for s in result_data.shape:
result_data.numel *= s
2024-04-28 18:43:07 -03:00
2024-05-01 02:12:47 -03:00
result_data.requires_grad = self.requires_grad or other.requires_grad
if result_data.requires_grad:
result_data.grad_fn = MatmulBackward(self, other)
2024-04-28 18:43:07 -03:00
return result_data
2024-04-26 18:38:07 -03:00
2024-05-06 01:21:18 -03:00
def __pow__(self, other):
2024-05-06 10:18:12 -03:00
other = float(other)
Tensor._C.tensor_pow_scalar.argtypes = [ctypes.POINTER(CTensor), ctypes.c_float]
Tensor._C.tensor_pow_scalar.restype = ctypes.POINTER(CTensor)
2024-05-06 01:21:18 -03:00
2024-05-06 10:18:12 -03:00
result_tensor_ptr = Tensor._C.tensor_pow_scalar(self.tensor, ctypes.c_float(other))
2024-05-06 01:21:18 -03:00
2024-05-06 10:18:12 -03:00
result_data = Tensor()
result_data.tensor = result_tensor_ptr
result_data.shape = self.shape.copy()
result_data.ndim = self.ndim
result_data.device = self.device
result_data.numel = self.numel
2024-05-06 10:18:12 -03:00
result_data.requires_grad = self.requires_grad
if result_data.requires_grad:
result_data.grad_fn = PowBackward(self, other)
2024-05-06 01:21:18 -03:00
return result_data
2024-05-06 10:18:12 -03:00
def __rpow__(self, other):
2024-05-06 01:21:18 -03:00
other = float(other)
2024-05-06 10:18:12 -03:00
Tensor._C.scalar_pow_tensor.argtypes = [ctypes.c_float, ctypes.POINTER(CTensor)]
Tensor._C.scalar_pow_tensor.restype = ctypes.POINTER(CTensor)
2024-05-06 01:21:18 -03:00
2024-05-06 10:18:12 -03:00
result_tensor_ptr = Tensor._C.scalar_pow_tensor(ctypes.c_float(other), self.tensor)
2024-05-06 01:21:18 -03:00
result_data = Tensor()
result_data.tensor = result_tensor_ptr
result_data.shape = self.shape.copy()
result_data.ndim = self.ndim
result_data.device = self.device
result_data.numel = self.numel
2024-05-06 01:21:18 -03:00
result_data.requires_grad = self.requires_grad
if result_data.requires_grad:
2024-05-06 10:18:12 -03:00
result_data.grad_fn = PowBackward(other, self)
return result_data
def __truediv__(self, other):
if isinstance(other, (int, float)):
other = float(other)
Tensor._C.tensor_div_scalar.argtypes = [ctypes.POINTER(CTensor), ctypes.c_float]
Tensor._C.tensor_div_scalar.restype = ctypes.POINTER(CTensor)
result_tensor_ptr = Tensor._C.tensor_div_scalar(self.tensor, ctypes.c_float(other))
result_data = Tensor()
result_data.tensor = result_tensor_ptr
result_data.shape = self.shape.copy()
result_data.ndim = self.ndim
result_data.device = self.device
result_data.numel = self.numel
2024-05-06 10:18:12 -03:00
result_data.requires_grad = self.requires_grad
if result_data.requires_grad:
result_data.grad_fn = DivisionBackward(self, other)
elif isinstance(self, Tensor) and isinstance(other, Tensor):
2024-05-21 01:22:38 -03:00
if other.numel == 1:
return self.__truediv__(other.tensor.contents.data[0])
2024-05-06 10:18:12 -03:00
Tensor._C.tensor_div_tensor.argtypes = [ctypes.POINTER(CTensor), ctypes.POINTER(CTensor)]
Tensor._C.tensor_div_tensor.restype = ctypes.POINTER(CTensor)
result_tensor_ptr = Tensor._C.tensor_div_tensor(self.tensor, other.tensor)
result_data = Tensor()
result_data.tensor = result_tensor_ptr
result_data.shape = self.shape.copy()
result_data.ndim = self.ndim
result_data.device = self.device
result_data.numel = self.numel
2024-05-06 10:18:12 -03:00
result_data.requires_grad = self.requires_grad or other.requires_grad
if result_data.requires_grad:
result_data.grad_fn = DivisionBackward(self, other)
2024-05-06 01:21:18 -03:00
return result_data
def __rtruediv__(self, other):
other = float(other)
Tensor._C.scalar_div_tensor.argtypes = [ctypes.c_float, ctypes.POINTER(CTensor)]
Tensor._C.scalar_div_tensor.restype = ctypes.POINTER(CTensor)
result_tensor_ptr = Tensor._C.scalar_div_tensor(ctypes.c_float(other), self.tensor)
2024-04-29 02:20:31 -03:00
2024-05-06 01:21:18 -03:00
result_data = Tensor()
result_data.tensor = result_tensor_ptr
result_data.shape = self.shape.copy()
result_data.ndim = self.ndim
result_data.device = self.device
result_data.numel = self.numel
2024-05-06 10:18:12 -03:00
2024-05-06 01:21:18 -03:00
result_data.requires_grad = self.requires_grad
if result_data.requires_grad:
result_data.grad_fn = DivisionBackward(other, self)
return result_data
2024-05-20 12:50:45 -03:00
def equal(self, other):
2024-05-20 15:24:55 -03:00
if isinstance(other, Tensor) and other.numel == 1:
# other is a single value tensor
other = self.zeros_like() + other
return self.equal(other)
2024-05-20 12:50:45 -03:00
if not isinstance(other, Tensor):
2024-05-20 15:24:55 -03:00
# other is a single value
if isinstance(other, (int, float)):
other = self.zeros_like() + other
return self.equal(other)
else:
return False
2024-05-20 12:50:45 -03:00
2024-05-20 16:42:45 -03:00
broadcasted_shape_add = []
2024-05-20 12:50:45 -03:00
2024-05-20 16:42:45 -03:00
# Function to determine if broadcasting is needed and get the broadcasted shape
def broadcast_shape(shape1, shape2):
if shape1 == shape2:
return shape1, False
max_len = max(len(shape1), len(shape2))
shape1 = [1] * (max_len - len(shape1)) + shape1
shape2 = [1] * (max_len - len(shape2)) + shape2
2024-05-20 12:50:45 -03:00
2024-05-20 16:42:45 -03:00
for dim1, dim2 in zip(shape1, shape2):
if dim1 != dim2 and dim1 != 1 and dim2 != 1:
raise ValueError("Shapes are not compatible for broadcasting")
broadcasted_shape_add.append(max(dim1, dim2))
return broadcasted_shape_add, True
broadcasted_shape_add, needs_broadcasting = broadcast_shape(self.shape, other.shape)
2024-05-20 12:50:45 -03:00
2024-05-20 16:42:45 -03:00
if needs_broadcasting:
if other.ndim == self.ndim - 1:
other = other.reshape([1] + other.shape)
elif self.ndim == other.ndim - 1:
self = self.reshape([1] + self.shape)
2024-05-20 20:16:41 -03:00
2024-05-20 16:42:45 -03:00
# Call equal_broadcasted_tensor if broadcasting is needed
Tensor._C.equal_broadcasted_tensor.argtypes = [ctypes.POINTER(CTensor), ctypes.POINTER(CTensor)]
Tensor._C.equal_broadcasted_tensor.restype = ctypes.POINTER(CTensor)
result_tensor_ptr = Tensor._C.equal_broadcasted_tensor(self.tensor, other.tensor)
result_data = Tensor()
result_data.tensor = result_tensor_ptr
result_data.shape = broadcasted_shape_add.copy()
result_data.ndim = len(broadcasted_shape_add)
result_data.device = self.device
result_data.numel = 1
for s in result_data.shape:
result_data.numel *= s
else:
Tensor._C.equal_tensor.argtypes = [ctypes.POINTER(CTensor), ctypes.POINTER(CTensor)]
Tensor._C.equal_tensor.restype = ctypes.POINTER(CTensor)
2024-05-20 12:50:45 -03:00
2024-05-20 16:42:45 -03:00
result_tensor_ptr = Tensor._C.equal_tensor(self.tensor, other.tensor)
result_data = Tensor()
result_data.tensor = result_tensor_ptr
result_data.shape = self.shape.copy()
result_data.ndim = self.ndim
result_data.device = self.device
result_data.numel = self.numel
return result_data
2024-05-20 12:50:45 -03:00
2024-05-06 01:21:18 -03:00
def log(self):
Tensor._C.log_tensor.argtypes = [ctypes.POINTER(CTensor)]
Tensor._C.log_tensor.restype = ctypes.POINTER(CTensor)
2024-04-29 02:20:31 -03:00
2024-05-06 01:21:18 -03:00
result_tensor_ptr = Tensor._C.log_tensor(self.tensor)
2024-04-29 02:20:31 -03:00
2024-04-30 01:16:13 -03:00
result_data = Tensor()
result_data.tensor = result_tensor_ptr
result_data.shape = self.shape.copy()
result_data.ndim = self.ndim
result_data.device = self.device
result_data.numel = self.numel
2024-04-30 01:16:13 -03:00
2024-05-01 00:49:07 -03:00
result_data.requires_grad = self.requires_grad
if result_data.requires_grad:
2024-05-06 01:21:18 -03:00
result_data.grad_fn = LogBackward(self)
2024-05-01 00:49:07 -03:00
2024-04-30 13:22:27 -03:00
return result_data
2024-05-20 20:16:41 -03:00
def sum(self, axis=None, keepdim=False):
2024-05-21 01:22:38 -03:00
if axis is not None and axis < 0:
axis = self.ndim + axis
2024-05-20 20:16:41 -03:00
if axis == None:
axis = -1
2024-05-21 01:22:38 -03:00
2024-05-22 19:55:00 -03:00
if axis > self.ndim - 1:
raise ValueError(f"Error: axis argument {axis} cannot be higher than tensor dimension {self.ndim}")
2024-05-17 19:08:09 -03:00
Tensor._C.sum_tensor.argtypes = [ctypes.POINTER(CTensor), ctypes.c_int, ctypes.c_bool]
2024-04-30 13:22:27 -03:00
Tensor._C.sum_tensor.restype = ctypes.POINTER(CTensor)
2024-05-17 19:08:09 -03:00
result_tensor_ptr = Tensor._C.sum_tensor(self.tensor, axis, keepdim)
2024-04-30 13:22:27 -03:00
result_data = Tensor()
result_data.tensor = result_tensor_ptr
2024-05-16 16:35:55 -03:00
if axis == -1:
2024-05-17 19:08:09 -03:00
if keepdim:
result_data.ndim = self.ndim
result_data.shape = [1] * self.ndim
else:
result_data.shape = [1]
result_data.ndim = 1
2024-05-16 16:35:55 -03:00
else:
2024-05-17 19:08:09 -03:00
if keepdim:
result_data.shape = self.shape[:axis] + [1] + self.shape[axis+1:]
else:
result_data.shape = self.shape[:axis] + self.shape[axis+1:]
2024-05-16 16:35:55 -03:00
result_data.ndim = len(result_data.shape)
2024-04-30 13:22:27 -03:00
result_data.device = self.device
result_data.numel = 1
2024-05-16 16:35:55 -03:00
for s in result_data.shape:
result_data.numel *= s
2024-04-30 13:22:27 -03:00
2024-04-30 20:34:58 -03:00
result_data.requires_grad = self.requires_grad
if result_data.requires_grad:
2024-05-17 19:08:09 -03:00
result_data.grad_fn = SumBackward(self, axis, keepdim=keepdim)
2024-05-20 12:50:45 -03:00
return result_data
2024-05-20 20:16:41 -03:00
def max(self, axis=None, keepdim=False):
2024-05-21 01:22:38 -03:00
if axis is not None and axis < 0:
axis = self.ndim + axis
2024-05-20 20:16:41 -03:00
if axis == None:
axis = -1
2024-05-21 01:22:38 -03:00
2024-05-22 19:55:00 -03:00
if axis > self.ndim - 1:
raise ValueError(f"Error: axis argument {axis} cannot be higher than tensor dimension {self.ndim}")
2024-05-20 12:50:45 -03:00
Tensor._C.max_tensor.argtypes = [ctypes.POINTER(CTensor), ctypes.c_int, ctypes.c_bool]
Tensor._C.max_tensor.restype = ctypes.POINTER(CTensor)
result_tensor_ptr = Tensor._C.max_tensor(self.tensor, axis, keepdim)
result_data = Tensor()
result_data.tensor = result_tensor_ptr
if axis == -1:
if keepdim:
result_data.ndim = self.ndim
result_data.shape = [1] * self.ndim
else:
result_data.shape = [1]
result_data.ndim = 1
else:
if keepdim:
result_data.shape = self.shape[:axis] + [1] + self.shape[axis+1:]
else:
result_data.shape = self.shape[:axis] + self.shape[axis+1:]
result_data.ndim = len(result_data.shape)
result_data.device = self.device
result_data.numel = 1
for s in result_data.shape:
result_data.numel *= s
result_data.requires_grad = self.requires_grad
if result_data.requires_grad:
result_data.grad_fn = MaxBackward(self, axis, keepdim=keepdim)
return result_data
2024-05-20 20:16:41 -03:00
def min(self, axis=None, keepdim=False):
2024-05-21 01:22:38 -03:00
if axis is not None and axis < 0:
axis = self.ndim + axis
2024-05-20 20:16:41 -03:00
if axis == None:
axis = -1
2024-05-22 19:55:00 -03:00
if axis > self.ndim - 1:
raise ValueError(f"Error: axis argument {axis} must be smaller than tensor dimension {self.ndim}")
2024-05-21 01:22:38 -03:00
2024-05-20 12:50:45 -03:00
Tensor._C.min_tensor.argtypes = [ctypes.POINTER(CTensor), ctypes.c_int, ctypes.c_bool]
Tensor._C.min_tensor.restype = ctypes.POINTER(CTensor)
result_tensor_ptr = Tensor._C.min_tensor(self.tensor, axis, keepdim)
result_data = Tensor()
result_data.tensor = result_tensor_ptr
if axis == -1:
if keepdim:
result_data.ndim = self.ndim
result_data.shape = [1] * self.ndim
else:
result_data.shape = [1]
result_data.ndim = 1
else:
if keepdim:
result_data.shape = self.shape[:axis] + [1] + self.shape[axis+1:]
else:
result_data.shape = self.shape[:axis] + self.shape[axis+1:]
result_data.ndim = len(result_data.shape)
result_data.device = self.device
result_data.numel = 1
for s in result_data.shape:
result_data.numel *= s
result_data.requires_grad = self.requires_grad
if result_data.requires_grad:
result_data.grad_fn = MinBackward(self, axis, keepdim=keepdim)
2024-04-30 20:09:08 -03:00
2024-05-01 02:12:47 -03:00
return result_data
2024-05-16 16:35:55 -03:00
2024-05-01 02:12:47 -03:00
def sin(self):
Tensor._C.sin_tensor.argtypes = [ctypes.POINTER(CTensor)]
Tensor._C.sin_tensor.restype = ctypes.POINTER(CTensor)
result_tensor_ptr = Tensor._C.sin_tensor(self.tensor)
result_data = Tensor()
result_data.tensor = result_tensor_ptr
result_data.shape = self.shape.copy()
result_data.ndim = self.ndim
result_data.device = self.device
result_data.numel = self.numel
result_data.requires_grad = self.requires_grad
if result_data.requires_grad:
result_data.grad_fn = SinBackward(self)
return result_data
def cos(self):
Tensor._C.cos_tensor.argtypes = [ctypes.POINTER(CTensor)]
Tensor._C.cos_tensor.restype = ctypes.POINTER(CTensor)
result_tensor_ptr = Tensor._C.cos_tensor(self.tensor)
result_data = Tensor()
result_data.tensor = result_tensor_ptr
result_data.shape = self.shape.copy()
result_data.ndim = self.ndim
result_data.device = self.device
result_data.numel = self.numel
result_data.requires_grad = self.requires_grad
if result_data.requires_grad:
result_data.grad_fn = CosBackward(self)
return result_data
2024-05-22 13:24:01 -03:00
def sigmoid(self):
Tensor._C.sigmoid_tensor.argtypes = [ctypes.POINTER(CTensor)]
Tensor._C.sigmoid_tensor.restype = ctypes.POINTER(CTensor)
result_tensor_ptr = Tensor._C.sigmoid_tensor(self.tensor)
result_data = Tensor()
result_data.tensor = result_tensor_ptr
result_data.shape = self.shape.copy()
result_data.ndim = self.ndim
result_data.device = self.device
result_data.numel = self.numel
result_data.requires_grad = self.requires_grad
if result_data.requires_grad:
result_data.grad_fn = SigmoidBackward(self)
return result_data
2024-05-16 19:11:15 -03:00
def transpose(self, axis1, axis2):
2024-05-04 13:37:18 -03:00
if axis1 < 0:
axis1 = self.ndim + axis1
if axis2 < 0:
axis2 = self.ndim + axis2
Tensor._C.transpose_axes_tensor.argtypes = [ctypes.POINTER(CTensor), ctypes.c_int, ctypes.c_int]
Tensor._C.transpose_axes_tensor.restype = ctypes.POINTER(CTensor)
result_tensor_ptr = Tensor._C.transpose_axes_tensor(self.tensor, axis1, axis2)
result_data = Tensor()
result_data.tensor = result_tensor_ptr
result_data.shape = self.shape.copy()
result_data.shape[axis1] = self.shape[axis2]
result_data.shape[axis2] = self.shape[axis1]
result_data.ndim = self.ndim
result_data.device = self.device
result_data.numel = self.numel
2024-05-04 13:37:18 -03:00
2024-05-04 13:47:20 -03:00
result_data.requires_grad = self.requires_grad
if result_data.requires_grad:
result_data.grad_fn = TransposeBackward(self, axis1, axis2)
2024-05-04 13:37:18 -03:00
return result_data
2024-05-01 02:12:47 -03:00
@property
def T(self):
Tensor._C.transpose_tensor.argtypes = [ctypes.POINTER(CTensor)]
Tensor._C.transpose_tensor.restype = ctypes.POINTER(CTensor)
result_tensor_ptr = Tensor._C.transpose_tensor(self.tensor)
result_data = Tensor()
result_data.tensor = result_tensor_ptr
2024-05-03 13:18:37 -03:00
result_data.shape = self.shape.copy()[::-1]
2024-05-02 21:14:26 -03:00
result_data.ndim = self.ndim
2024-05-01 02:12:47 -03:00
result_data.device = self.device
result_data.numel = self.numel
result_data.requires_grad = self.requires_grad
2024-05-07 12:48:06 -03:00
if result_data.requires_grad:
result_data.grad_fn = TBackward(self)
return result_data
def detach(self):
self.grad = None
self.grad_fn = None
2024-05-01 02:12:47 -03:00
2024-05-07 11:52:56 -03:00
return self