| """
|
| Tensor operations for distributed computing.
|
| """
|
| import torch
|
| import numpy as np
|
| from typing import Dict, List, Optional, Union, Tuple
|
|
|
| class TensorOps:
|
| """Utility class for distributed tensor operations."""
|
|
|
| @staticmethod
|
| def split_tensor(tensor: torch.Tensor, num_parts: int) -> List[torch.Tensor]:
|
| """Split a tensor into multiple parts for distributed processing."""
|
| return torch.chunk(tensor, num_parts)
|
|
|
| @staticmethod
|
| def merge_tensors(tensors: List[torch.Tensor], dim: int = 0) -> torch.Tensor:
|
| """Merge multiple tensors back into a single tensor."""
|
| return torch.cat(tensors, dim=dim)
|
|
|
| @staticmethod
|
| def average_gradients(gradients: List[Dict[str, torch.Tensor]]) -> Dict[str, torch.Tensor]:
|
| """Average gradients from multiple workers."""
|
| avg_gradients = {}
|
| for key in gradients[0].keys():
|
| avg_gradients[key] = torch.mean(torch.stack([g[key] for g in gradients]), dim=0)
|
| return avg_gradients
|
|
|
| @staticmethod
|
| def serialize_tensor(tensor: torch.Tensor) -> Dict[str, Union[List, str]]:
|
| """Serialize a tensor for storage/transmission."""
|
| return {
|
| 'data': tensor.cpu().numpy().tolist(),
|
| 'shape': list(tensor.shape),
|
| 'dtype': str(tensor.dtype)
|
| }
|
|
|
| @staticmethod
|
| def deserialize_tensor(tensor_dict: Dict[str, Union[List, str]]) -> torch.Tensor:
|
| """Deserialize a tensor from storage/transmission format."""
|
| data = np.array(tensor_dict['data'])
|
| shape = tensor_dict['shape']
|
| dtype = getattr(torch, tensor_dict['dtype'].split('.')[-1])
|
| return torch.tensor(data, dtype=dtype).reshape(shape)
|
|
|
| @staticmethod
|
| def gradient_clipping(gradients: Dict[str, torch.Tensor], max_norm: float) -> Dict[str, torch.Tensor]:
|
| """Apply gradient clipping to prevent exploding gradients."""
|
| for k, v in gradients.items():
|
| if v is not None:
|
| torch.nn.utils.clip_grad_norm_(v, max_norm)
|
| return gradients
|
|
|
| @staticmethod
|
| def reduce_precision(tensor: torch.Tensor, bits: int = 16) -> torch.Tensor:
|
| """Reduce tensor precision for efficient transmission."""
|
| if bits == 16:
|
| return tensor.half()
|
| elif bits == 32:
|
| return tensor.float()
|
| else:
|
| raise ValueError("Unsupported precision bits")
|
|
|
| @staticmethod
|
| def shard_tensor(tensor: torch.Tensor, shard_size: int) -> List[torch.Tensor]:
|
| """Shard a tensor into smaller pieces for distributed processing."""
|
| return [tensor[i:i + shard_size] for i in range(0, tensor.size(0), shard_size)]
|
|
|
| @staticmethod
|
| def compute_parameter_norm(parameters: Dict[str, torch.Tensor]) -> float:
|
| """Compute the total norm of all parameters."""
|
| total_norm = 0.0
|
| for param in parameters.values():
|
| total_norm += param.norm().item() ** 2
|
| return total_norm ** 0.5 |
