File size: 7,765 Bytes
7a0c684 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 |
"""
Advanced tensor storage operations extending the base storage system.
"""
import numpy as np
from typing import Dict, Any, Optional, Union, List, Tuple
import threading
from http_storage import LocalStorage
import logging
import json
class TensorOps:
"""Tensor operations for storage system"""
@staticmethod
def serialize_tensor(tensor: np.ndarray) -> Tuple[bytes, Dict[str, Any]]:
"""Serialize tensor to bytes with metadata"""
metadata = {
'shape': tensor.shape,
'dtype': str(tensor.dtype),
'strides': tensor.strides
}
return tensor.tobytes(), metadata
@staticmethod
def deserialize_tensor(data: bytes, metadata: Dict[str, Any]) -> np.ndarray:
"""Deserialize tensor from bytes and metadata"""
tensor = np.frombuffer(data, dtype=np.dtype(metadata['dtype']))
return tensor.reshape(metadata['shape'])
class TensorStorage(LocalStorage):
"""
Enhanced storage implementation with tensor operations.
Extends LocalStorage with advanced tensor manipulation capabilities.
"""
def __init__(self, db_url: str = None):
super().__init__(db_url)
self._ops = TensorOps()
def store_tensor(self, tensor_id: str, tensor: np.ndarray, metadata: Optional[Dict] = None) -> bool:
"""Store a tensor with its metadata"""
try:
# Serialize tensor
data, tensor_meta = self._ops.serialize_tensor(tensor)
# Merge with additional metadata
if metadata:
tensor_meta.update(metadata)
# Store in database
success = self._store_in_db('tensors', tensor_id, data, tensor_meta)
if success:
# Update stats
with self._lock:
self.stats['tensor_count'] += 1
self.stats['total_size'] += len(data)
return success
except Exception as e:
logging.error(f"Error storing tensor {tensor_id}: {str(e)}")
return False
def get_tensor(self, tensor_id: str) -> Optional[np.ndarray]:
"""Retrieve a tensor by ID"""
try:
# Get from database
result = self.conn.execute("""
SELECT data, metadata
FROM tensors
WHERE id = ?
""", [tensor_id]).fetchone()
if not result:
return None
data, metadata = result
metadata = json.loads(metadata) if isinstance(metadata, str) else metadata
# Deserialize tensor
return self._ops.deserialize_tensor(data, metadata)
except Exception as e:
logging.error(f"Error retrieving tensor {tensor_id}: {str(e)}")
return None
def matmul(self, tensor_id_a: str, tensor_id_b: str, result_id: str) -> Optional[str]:
"""Perform matrix multiplication between two stored tensors"""
tensor_a = self.get_tensor(tensor_id_a)
tensor_b = self.get_tensor(tensor_id_b)
if tensor_a is None or tensor_b is None:
return None
try:
result = np.matmul(tensor_a, tensor_b)
if self.store_tensor(result_id, result):
return result_id
return None
except Exception as e:
logging.error(f"Error in matrix multiplication: {str(e)}")
return None
def add(self, tensor_id_a: str, tensor_id_b: str, result_id: str) -> Optional[str]:
"""Add two stored tensors"""
tensor_a = self.get_tensor(tensor_id_a)
tensor_b = self.get_tensor(tensor_id_b)
if tensor_a is None or tensor_b is None:
return None
try:
result = tensor_a + tensor_b
if self.store_tensor(result_id, result):
return result_id
return None
except Exception as e:
logging.error(f"Error in tensor addition: {str(e)}")
return None
def multiply(self, tensor_id_a: str, tensor_id_b: str, result_id: str) -> Optional[str]:
"""Element-wise multiply two stored tensors"""
tensor_a = self.get_tensor(tensor_id_a)
tensor_b = self.get_tensor(tensor_id_b)
if tensor_a is None or tensor_b is None:
return None
try:
result = tensor_a * tensor_b
if self.store_tensor(result_id, result):
return result_id
return None
except Exception as e:
logging.error(f"Error in tensor multiplication: {str(e)}")
return None
def transpose(self, tensor_id: str, result_id: str) -> Optional[str]:
"""Transpose a stored tensor"""
tensor = self.get_tensor(tensor_id)
if tensor is None:
return None
try:
result = np.transpose(tensor)
if self.store_tensor(result_id, result):
return result_id
return None
except Exception as e:
logging.error(f"Error in tensor transpose: {str(e)}")
return None
def reshape(self, tensor_id: str, new_shape: Tuple[int, ...], result_id: str) -> Optional[str]:
"""Reshape a stored tensor"""
tensor = self.get_tensor(tensor_id)
if tensor is None:
return None
try:
result = tensor.reshape(new_shape)
if self.store_tensor(result_id, result):
return result_id
return None
except Exception as e:
logging.error(f"Error in tensor reshape: {str(e)}")
return None
def split(self, tensor_id: str, indices_or_sections: Union[int, List[int]], axis: int = 0) -> List[str]:
"""Split a stored tensor into multiple tensors"""
tensor = self.get_tensor(tensor_id)
if tensor is None:
return []
try:
# Split the tensor
split_results = np.split(tensor, indices_or_sections, axis=axis)
# Store each split result
result_ids = []
for i, split_tensor in enumerate(split_results):
result_id = f"{tensor_id}_split_{i}"
if self.store_tensor(result_id, split_tensor):
result_ids.append(result_id)
return result_ids
except Exception as e:
logging.error(f"Error in tensor split: {str(e)}")
return []
def concatenate(self, tensor_ids: List[str], result_id: str, axis: int = 0) -> Optional[str]:
"""Concatenate multiple stored tensors"""
tensors = [self.get_tensor(tid) for tid in tensor_ids]
if None in tensors:
return None
try:
result = np.concatenate(tensors, axis=axis)
if self.store_tensor(result_id, result):
return result_id
return None
except Exception as e:
logging.error(f"Error in tensor concatenation: {str(e)}")
return None
|