Spaces:
Sleeping
Sleeping
File size: 3,343 Bytes
43464e3 |
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 |
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())
|