Spaces:
Sleeping
Sleeping
| import asyncio | |
| import websockets | |
| import json | |
| import numpy as np | |
| from typing import List, Any, Optional, Dict | |
| class TensorCoreArray: | |
| def __init__(self, num_tensor_cores: int, bits: int, bandwidth_tbps: float): | |
| self.num_tensor_cores = num_tensor_cores | |
| self.bits = bits | |
| self.bandwidth_tbps = bandwidth_tbps | |
| self.initialized = False | |
| def initialize(self): | |
| print(f"Initializing {self.num_tensor_cores} tensor cores with {self.bits}-bit precision...") | |
| self.initialized = True | |
| def matmul(self, matrix_a: List[List[float]], matrix_b: List[List[float]]) -> List[List[float]]: | |
| if not self.initialized: | |
| raise RuntimeError("Tensor cores not initialized. Call initialize() first.") | |
| np_a = np.array(matrix_a) | |
| np_b = np.array(matrix_b) | |
| if np_a.shape[1] != np_b.shape[0]: | |
| raise ValueError("Matrix dimensions incompatible for multiplication") | |
| result = np.matmul(np_a, np_b) | |
| return result.tolist() | |
| async def send_tensor_data(self, uri: str, tensor_id: str, data: np.ndarray): | |
| async with websockets.connect(uri) as websocket: | |
| payload = { | |
| "operation": "tensor_data", | |
| "type": "send", | |
| "tensor_id": tensor_id, | |
| "data": data.tolist() | |
| } | |
| await websocket.send(json.dumps(payload)) | |
| response = await websocket.recv() | |
| return json.loads(response) | |
| async def receive_tensor_data(self, uri: str, tensor_id: str) -> Optional[np.ndarray]: | |
| async with websockets.connect(uri) as websocket: | |
| payload = { | |
| "operation": "tensor_data", | |
| "type": "receive", | |
| "tensor_id": tensor_id | |
| } | |
| await websocket.send(json.dumps(payload)) | |
| response = await websocket.recv() | |
| response_data = json.loads(response) | |
| if response_data.get("status") == "success": | |
| return np.array(response_data["data"]) | |
| return None | |
| def get_status(self) -> Dict[str, Any]: | |
| return { | |
| "num_tensor_cores": self.num_tensor_cores, | |
| "bits": self.bits, | |
| "bandwidth_tbps": self.bandwidth_tbps, | |
| "initialized": self.initialized | |
| } | |
| if __name__ == "__main__": | |
| async def test_tensor_core_array(): | |
| tca = TensorCoreArray(num_tensor_cores=10, bits=32, bandwidth_tbps=1.0) | |
| tca.initialize() | |
| matrix_a = [[1, 2], [3, 4]] | |
| matrix_b = [[5, 6], [7, 8]] | |
| result = tca.matmul(matrix_a, matrix_b) | |
| print(f"Matrix multiplication result: {result}") | |
| # Example of sending/receiving tensor data (requires a running WebSocket server) | |
| # uri = "ws://localhost:7860/ws" | |
| # tensor_id = "test_tensor" | |
| # dummy_data = np.array([[10, 20], [30, 40]]) | |
| # | |
| # print(f"Sending tensor data: {dummy_data.tolist()}") | |
| # send_response = await tca.send_tensor_data(uri, tensor_id, dummy_data) | |
| # print(f"Send response: {send_response}") | |
| # | |
| # received_data = await tca.receive_tensor_data(uri, tensor_id) | |
| # print(f"Received tensor data: {received_data.tolist() if received_data is not None else None}") | |
| asyncio.run(test_tensor_core_array()) | |