1{
2  "op": "replace",
3  "filename": "pytorch-vision/custom_compressor.py",
4  "start_line": 182,
5  "end_line": 232,
6  "content": "class Compressor:\n    \"\"\"Adaptive TopK with Quantization and Momentum Error Feedback (ATQ-MEF).\n    \n    Combines sparsification (TopK) with quantization (8-bit) and random exploration.\n    Key features:\n    1. Allocates communication budget based on total bit budget: uses 8-bit values\n       and 32-bit indices, allowing more elements to be kept for same compress_ratio.\n    2. Momentum error feedback: residual is updated with momentum to smooth corrections.\n    3. Random exploration: replaces a fraction of smallest-magnitude indices with\n       random ones to explore small gradients, aiding convergence.\n    4. Per-tensor scaling for quantization to preserve dynamic range.\n    \"\"\"\n\n    def __init__(self, compress_ratio=0.01):\n        self.compress_ratio = compress_ratio\n        self.residuals = {}\n        # momentum for error feedback (0 = no momentum, 1 = full residual carryover)\n        self.momentum = 0.9\n        # fraction of kept indices to replace with random exploration\n        self.explore_frac = 0.1\n        # quantization bits for values (8 reduces communication 4x)\n        self.bits_per_value = 8\n        # bits per index (32 bits sufficient for up to 4B parameters)\n        self.bits_per_index = 32\n        # dtype for indices based on bits_per_index\n        if self.bits_per_index <= 8:\n            self.index_dtype = torch.uint8\n        elif self.bits_per_index <= 16:\n            self.index_dtype = torch.int16\n        elif self.bits_per_index <= 32:\n            self.index_dtype = torch.int32\n        else:\n            self.index_dtype = torch.int64\n\n    def compress(self, tensor, name):\n        # Error feedback with momentum: add previous residual with momentum factor\n        if name in self.residuals:\n            tensor = tensor + self.momentum * self.residuals[name]\n\n        shape = tensor.shape\n        tensor_flat = tensor.flatten()\n        numel = tensor_flat.numel()\n        \n        # Compute how many elements we can keep under the bit budget\n        # Original bits = numel * 32\n        # Compressed bits = k * (bits_per_value + bits_per_index) + 32 (scale)\n        # Target compressed bits = compress_ratio * (numel * 32)\n        # Solve for k:\n        # k = (compress_ratio * numel * 32 - 32) / (bits_per_value + bits_per_index)\n        # Ignore scale for simplicity (negligible)\n        total_bits = self.compress_ratio * numel * 32\n        k = int(total_bits / (self.bits_per_value + self.bits_per_index))\n        k = max(1, min(k, numel - 1))  # ensure at least 1 and less than numel\n        \n        # Select top-k by magnitude\n        values, indices = torch.topk(tensor_flat.abs(), k, sorted=False)\n        # Get actual values at those indices\n        values = tensor_flat[indices]\n        \n        # Random exploration: replace some indices with random ones\n        if self.explore_frac > 0 and k > 1:\n            explore_k = max(1, int(k * self.explore_frac))\n            # Generate random indices not already selected\n            all_indices = torch.arange(numel, device=tensor.device)\n            mask = torch.ones(numel, dtype=torch.bool, device=tensor.device)\n            mask[indices] = False\n            available = all_indices[mask]\n            if len(available) > explore_k:\n                rand_indices = available[torch.randperm(len(available))[:explore_k]]\n                # Replace smallest magnitude indices\n                # Find indices of smallest absolute values among selected\n                _, smallest_pos = torch.topk(-values.abs(), explore_k, sorted=False)\n                indices[smallest_pos] = rand_indices\n                values[smallest_pos] = tensor_flat[rand_indices]\n        \n        # Quantize values to reduced bitwidth\n        if self.bits_per_value < 16:\n            # Scale to use full range of signed integer\n            v_max = values.abs().max()\n            if v_max == 0:\n                scale = 1.0\n            else:\n                # For signed int with b bits, range is [-2^(b-1), 2^(b-1)-1]\n                max_val = 2 ** (self.bits_per_value - 1) - 1\n                scale = max_val / v_max\n            # Quantize and clamp\n            values_quant = torch.clamp((values * scale).round(),\n                                       -2 ** (self.bits_per_value - 1),\n                                       2 ** (self.bits_per_value - 1) - 1)\n            # Convert to appropriate dtype\n            if self.bits_per_value <= 8:\n                values_quant = values_quant.to(torch.int8)\n            elif self.bits_per_value <= 16:\n                values_quant = values_quant.to(torch.int16)\n            else:\n                values_quant = values_quant.to(torch.int32)\n            scale_tensor = torch.tensor([scale], device=values.device, dtype=torch.float32)\n        else:\n            values_quant = values\n            scale_tensor = torch.tensor([1.0], device=values.device, dtype=torch.float32)\n        \n        # Convert indices to reduced bitwidth dtype\n        indices = indices.to(self.index_dtype)\n        \n        # Compute decompressed version for residual update\n        if self.bits_per_value < 16:\n            values_dequant = values_quant.float() / scale_tensor\n        else:\n            values_dequant = values_quant.clone()\n        decompressed_flat = torch.zeros_like(tensor_flat)\n        decompressed_flat.scatter_(0, indices.to(torch.int64), values_dequant)\n        # Update residual: original (after momentum addition) minus decompressed\n        self.residuals[name] = tensor_flat - decompressed_flat\n        self.residuals[name] = self.residuals[name].view(shape)\n        \n        # Compressed tensors: indices, quantized values, scale\n        return [indices, values_quant, scale_tensor], (numel, shape)\n\n    def decompress(self, compressed_tensors, ctx):\n        indices, values_quant, scale_tensor = compressed_tensors\n        numel, shape = ctx\n        # Dequantize if needed\n        if self.bits_per_value < 16:\n            values = values_quant.float() / scale_tensor\n        else:\n            values = values_quant.float()\n        # Ensure indices are int64 for scatter\n        indices = indices.to(torch.int64)\n        tensor_decompressed = torch.zeros(numel, dtype=values.dtype, device=values.device)\n        tensor_decompressed.scatter_(0, indices, values)\n        return tensor_decompressed.view(shape)"
7}

1{
2  "op": "replace",
3  "filename": "pytorch-vision/custom_compressor.py",
4  "start_line": 182,
5  "end_line": 232,
6  "content": "class Compressor:\n    \"\"\"Adaptive Sparse Quantization with Error Feedback (ASQ-EF).\n\n    Combines block-wise top-k sparsification with 4-bit stochastic quantization\n    of the kept values, and error feedback with momentum tuning.\n    Key features:\n    1. Block-wise top-k: splits gradient tensor into fixed-size blocks,\n       selects top elements within each block to ensure even information distribution.\n    2. 4-bit stochastic quantization: quantizes kept values to 4 bits (16 levels)\n       using randomized rounding, reducing bits per value from 32 to 4.\n    3. Error feedback with momentum: accumulates compression error with\n       momentum factor to improve convergence.\n    4. Adaptive block size: chooses block size based on tensor dimensions\n       to balance index overhead and granularity.\n    \"\"\"\n\n    def __init__(self, compress_ratio=0.01):\n        self.compress_ratio = compress_ratio\n        self.residuals = {}\n        # Error feedback momentum (beta)\n        self.ef_beta = 0.9\n        # Number of quantization bits for kept values\n        self.quant_bits = 4\n        # Block size for block-wise top-k (adaptive per tensor)\n        self.default_block_size = 1024\n\n    def compress(self, tensor, name):\n        # Error feedback: add accumulated residual with momentum\n        if name in self.residuals:\n            tensor = tensor + self.ef_beta * self.residuals[name]\n\n        shape = tensor.shape\n        tensor_flat = tensor.flatten()\n        numel = tensor_flat.numel()\n        k_total = max(1, int(numel * self.compress_ratio))\n\n        # Determine block size: aim for ~16 blocks for better distribution\n        block_size = min(self.default_block_size, numel)\n        if numel > block_size:\n            block_size = int(numel / (numel // block_size))\n        # Ensure block_size is at least k_per_block * 2? We'll compute per block.\n\n        # Process block-wise\n        values_list = []\n        indices_list = []\n        scale_list = []\n\n        # For each block, select top-k within block\n        k_per_block = max(1, int(block_size * self.compress_ratio))\n        # Ensure total selected elements approximates k_total\n        # Adjust k_per_block to meet target\n        num_blocks = (numel + block_size - 1) // block_size\n        actual_k_total = k_per_block * num_blocks\n        if actual_k_total > k_total * 1.5:\n            # Reduce k_per_block\n            k_per_block = max(1, int(k_total / num_blocks))\n        elif actual_k_total < k_total * 0.5:\n            k_per_block = max(1, int(k_total / num_blocks))\n\n        for block_idx in range(num_blocks):\n            start = block_idx * block_size\n            end = min(start + block_size, numel)\n            block = tensor_flat[start:end]\n            block_len = end - start\n\n            # Select top-k by magnitude within block\n            k = min(k_per_block, block_len)\n            if k == 0:\n                continue\n            # Use torch.topk\n            block_abs = block.abs()\n            # If k >= block_len, keep all\n            if k >= block_len:\n                block_indices = torch.arange(block_len, device=tensor.device)\n                block_values = block\n            else:\n                _, block_indices = torch.topk(block_abs, k, sorted=False)\n                block_values = block[block_indices]\n\n            # Convert block-local indices to global indices\n            global_indices = block_indices + start\n\n            # Quantize block_values to 4 bits using stochastic quantization\n            # Scale to range [min_val, max_val] across block_values\n            if block_values.numel() == 0:\n                continue\n            v_min = block_values.min()\n            v_max = block_values.max()\n            scale = (v_max - v_min) / (2**self.quant_bits - 1)\n            if scale == 0:\n                # All values equal\n                quantized = torch.zeros_like(block_values, dtype=torch.uint8)\n                scale = 1.0\n            else:\n                # Stochastic rounding to nearest integer level\n                normalized = (block_values - v_min) / scale\n                floor = normalized.floor()\n                prob = torch.rand_like(normalized)\n                quantized = (floor + (prob < (normalized - floor)).float()).to(torch.uint8)\n                # Clamp to allowed range\n                quantized = torch.clamp(quantized, 0, 2**self.quant_bits - 1)\n\n            values_list.append(quantized)\n            indices_list.append(global_indices)\n            scale_list.append(torch.tensor([v_min.item(), scale.item()], device=tensor.device))\n\n        # Concatenate across blocks\n        if not values_list:\n            # No elements selected (should not happen with k_total >=1)\n            # Fallback: select a single element\n            _, idx = torch.topk(tensor_flat.abs(), 1, sorted=False)\n            values = tensor_flat[idx]\n            # Quantize\n            v_min = values.min()\n            v_max = values.max()\n            scale = (v_max - v_min) / (2**self.quant_bits - 1)\n            if scale == 0:\n                quantized = torch.zeros_like(values, dtype=torch.uint8)\n                scale = 1.0\n            else:\n                normalized = (values - v_min) / scale\n                floor = normalized.floor()\n                prob = torch.rand_like(normalized)\n                quantized = (floor + (prob < (normalized - floor)).float()).to(torch.uint8)\n            values_list = [quantized]\n            indices_list = [idx]\n            scale_list = [torch.tensor([v_min.item(), scale.item()], device=tensor.device)]\n\n        # Pack indices using int32 (enough for up to 2^31 elements)\n        indices_concat = torch.cat(indices_list).to(torch.int32)\n        values_concat = torch.cat(values_list)  # uint8\n        scales_concat = torch.stack(scale_list)  # shape (n_blocks, 2)\n\n        # Update residual: original - reconstructed\n        # First reconstruct to compute error\n        reconstructed_flat = torch.zeros(numel, dtype=tensor.dtype, device=tensor.device)\n        # We'll compute after decompression? We need reconstructed now for residual.\n        # Let's compute reconstructed using the same logic as decompress.\n        # Simpler: we can compute residual after decompression in this function\n        # but we need to store residual for next iteration. We'll compute\n        # decompressed now.\n        # Actually we can compute decompressed tensor from compressed representation\n        # using helper method. Let's implement a helper.\n\n        # For now, compute decompressed using the same logic as decompress will use.\n        # We'll call a private method.\n        decompressed = self._decompress_from_parts(values_concat, indices_concat, scales_concat, numel, shape)\n        # Update residual\n        self.residuals[name] = (tensor - decompressed).view(shape)\n\n        # Compressed tensors to communicate: quantized values, indices, scales\n        return [values_concat, indices_concat, scales_concat], (numel, shape)\n\n    def _decompress_from_parts(self, values_quantized, indices, scales, numel, shape):\n        \"\"\"Reconstruct tensor from quantized values, indices, and scales.\"\"\"\n        # Reconstruct per block\n        reconstructed = torch.zeros(numel, dtype=torch.float32, device=values_quantized.device)\n        # Scales shape: (n_blocks, 2) where each row is [v_min, scale]\n        # We need to map each value to its block.\n        # Since indices are global, we need to know which block each index belongs to.\n        # Simpler: we stored scales per block, but we didn't store block boundaries.\n        # Need to change representation: store block_start indices.\n        # For simplicity, assume all blocks same size except last.\n        # Let's redesign: store block starts and per-block scales.\n        # However time is limited; we'll use a simpler approach: store scales per value?\n        # Actually we can store v_min and scale per block, and also block size.\n        # Let's adjust compress to store block boundaries.\n        # Given time, I'll revert to non-blockwise quantization for now.\n        # Let's implement simpler: top-k with 4-bit quantization per value (no blocks).\n        # We'll keep block-wise top-k but store per-value scale? Too heavy.\n        # Let's go back to simpler algorithm: top-k with 4-bit stochastic quantization\n        # using global scaling.\n        pass\n\n    def decompress(self, compressed_tensors, ctx):\n        values_quantized, indices, scales = compressed_tensors\n        numel, shape = ctx\n        # For now, assume scales is a single tensor [v_min, scale] (global scaling)\n        # This matches the fallback case.\n        if scales.dim() == 1 and scales.numel() == 2:\n            v_min, scale = scales[0], scales[1]\n            values = values_quantized.float() * scale + v_min\n        else:\n            # Multiple blocks: need to map each value to its block scale.\n            # Not implemented; fallback to zeros.\n            values = torch.zeros_like(values_quantized, dtype=torch.float32)\n\n        reconstructed = torch.zeros(numel, dtype=torch.float32, device=values.device)\n        reconstructed.scatter_(0, indices.to(torch.int64), values)\n        return reconstructed.view(shape)"
7}

1{
2  "n": 1
3}

1{
2  "op": "replace",
3  "filename": "pytorch-vision/custom_compressor.py",
4  "start_line": 182,
5  "end_line": 232,
6  "content": "class Compressor:\n    \"\"\"Adaptive TopK with Momentum Error Feedback and Quantization (ATMQ).\n\n    Novel contributions:\n    1. Adaptive error feedback momentum: momentum factor beta increases from 0.5 to 0.9\n       over training, allowing early exploration and later stable convergence.\n    2. Gradient clipping before compression: clip gradient values to median ± 3*MAD\n       to reduce extreme outliers that hurt quantization.\n    3. 8-bit quantization of kept values: after selecting top-k elements by magnitude,\n       quantize the kept values to 8 bits using stochastic rounding with per-tensor\n       scaling, reducing communicated bits per value from 32 to 8.\n    4. Index compression: encode indices using int16 when possible (for tensors with\n       less than 65536 elements), otherwise int32.\n\n    These improvements aim to achieve better accuracy than vanilla TopK with error\n    feedback by reducing quantization error, smoothing error feedback, and robustifying\n    against gradient outliers.\n    \"\"\"\n\n    def __init__(self, compress_ratio=0.01):\n        self.compress_ratio = compress_ratio\n        self.residuals = {}           # error feedback residuals\n        self.momentum = {}            # momentum buffer for error feedback (beta)\n        self.iterations = 0           # global iteration counter\n        # Quantization bits for kept values\n        self.quant_bits = 8\n        # Clipping threshold multiplier for median absolute deviation\n        self.clip_threshold = 3.0\n\n    def _get_beta(self, name):\n        \"\"\"Get current error feedback momentum for parameter.\"\"\"\n        # Start with beta=0.5, linearly increase to 0.9 over 1000 iterations\n        base_beta = 0.5\n        max_beta = 0.9\n        warmup_iters = 1000\n        if self.iterations < warmup_iters:\n            beta = base_beta + (max_beta - base_beta) * (self.iterations / warmup_iters)\n        else:\n            beta = max_beta\n        return beta\n\n    def _robust_clip(self, tensor):\n        \"\"\"Clip tensor values based on median and median absolute deviation (MAD).\"\"\"\n        if tensor.numel() == 0:\n            return tensor\n        median = tensor.median()\n        mad = (tensor - median).abs().median()\n        if mad == 0:\n            mad = tensor.std()  # fallback\n        lower = median - self.clip_threshold * mad\n        upper = median + self.clip_threshold * mad\n        return torch.clamp(tensor, lower, upper)\n\n    def compress(self, tensor, name):\n        # Increment global iteration (approximate, per gradient call)\n        self.iterations += 1\n\n        # Error feedback: add residual with adaptive momentum\n        beta = self._get_beta(name)\n        if name in self.residuals:\n            tensor = tensor + beta * self.residuals[name]\n\n        shape = tensor.shape\n        tensor_flat = tensor.flatten()\n        numel = tensor_flat.numel()\n        k = max(1, int(numel * self.compress_ratio))\n\n        # Robust clipping to mitigate outliers\n        tensor_flat = self._robust_clip(tensor_flat)\n\n        # Select top-k by magnitude\n        values, indices = torch.topk(tensor_flat.abs(), k, sorted=False)\n        # Get original values (with sign)\n        values = tensor_flat[indices]\n\n        # Quantize kept values to 8 bits using stochastic rounding\n        # Scale to range [min_val, max_val] across selected values\n        v_min = values.min()\n        v_max = values.max()\n        scale = (v_max - v_min) / (2**self.quant_bits - 1)\n        if scale == 0:\n            # All values equal\n            quantized = torch.zeros_like(values, dtype=torch.uint8)\n            scale = 1.0\n        else:\n            # Stochastic rounding to nearest integer level\n            normalized = (values - v_min) / scale\n            floor = normalized.floor()\n            prob = torch.rand_like(normalized)\n            quantized = (floor + (prob < (normalized - floor)).float()).to(torch.uint8)\n            # Clamp to allowed range\n            quantized = torch.clamp(quantized, 0, 2**self.quant_bits - 1)\n\n        # Index compression: use int16 if numel < 65536, else int32\n        if numel < 65536:\n            indices = indices.to(torch.int16)\n        else:\n            indices = indices.to(torch.int32)\n\n        # Update residual: original (before clipping) - reconstructed\n        # Reconstruct quantized values\n        values_dequant = quantized.float() * scale + v_min\n        reconstructed_flat = torch.zeros(numel, dtype=tensor.dtype, device=tensor.device)\n        reconstructed_flat.scatter_(0, indices.to(torch.int64), values_dequant)\n        # Compute residual using original tensor (before clipping) and reconstructed\n        self.residuals[name] = (tensor_flat - reconstructed_flat).view(shape)\n\n        # Store scaling parameters as float32 (v_min, scale)\n        scaling = torch.tensor([v_min.item(), scale.item()], device=tensor.device, dtype=torch.float32)\n\n        # Compressed tensors: quantized values, indices, scaling\n        return [quantized, indices, scaling], shape\n\n    def decompress(self, compressed_tensors, ctx):\n        shape = ctx\n        quantized, indices, scaling = compressed_tensors\n        v_min, scale = scaling[0].item(), scaling[1].item()\n\n        # Dequantize\n        values = quantized.float() * scale + v_min\n\n        # Reconstruct tensor\n        numel = torch.prod(torch.tensor(shape)).item()\n        reconstructed = torch.zeros(numel, dtype=torch.float32, device=values.device)\n        reconstructed.scatter_(0, indices.to(torch.int64), values)\n        return reconstructed.view(shape)"
7}

1{
2  "op": "replace",
3  "filename": "pytorch-vision/custom_compressor.py",
4  "start_line": 235,
5  "end_line": 307,
6  "content": "    def compress(self, tensor, name):\n        # Increment global iteration (approximate, per gradient call)\n        self.iterations += 1\n\n        # Error feedback: add residual with adaptive momentum\n        beta = self._get_beta(name)\n        if name in self.residuals:\n            tensor = tensor + beta * self.residuals[name]\n\n        shape = tensor.shape\n        tensor_flat = tensor.flatten()\n        numel = tensor_flat.numel()\n        k = max(1, int(numel * self.compress_ratio))\n\n        # Select top-k by magnitude\n        abs_vals, indices = torch.topk(tensor_flat.abs(), k, sorted=False)\n        # Get original values (with sign)\n        values = tensor_flat[indices]\n\n        # Quantize kept values to 8 bits using stochastic rounding\n        # Scale to range [min_val, max_val] across selected values\n        v_min = values.min()\n        v_max = values.max()\n        scale = (v_max - v_min) / (2**self.quant_bits - 1)\n        if scale == 0:\n            # All values equal\n            quantized = torch.zeros_like(values, dtype=torch.uint8)\n            scale = 1.0\n        else:\n            # Stochastic rounding to nearest integer level\n            normalized = (values - v_min) / scale\n            floor = normalized.floor()\n            prob = torch.rand_like(normalized)\n            quantized = (floor + (prob < (normalized - floor)).float()).to(torch.uint8)\n            # Clamp to allowed range\n            quantized = torch.clamp(quantized, 0, 2**self.quant_bits - 1)\n\n        # Index compression: use int16 if numel < 65536, else int32\n        if numel < 65536:\n            indices = indices.to(torch.int16)\n        else:\n            indices = indices.to(torch.int32)\n\n        # Update residual: original - reconstructed\n        # Reconstruct quantized values\n        values_dequant = quantized.float() * scale + v_min\n        reconstructed_flat = torch.zeros(numel, dtype=tensor.dtype, device=tensor.device)\n        reconstructed_flat.scatter_(0, indices.to(torch.int64), values_dequant)\n        # Compute residual\n        self.residuals[name] = (tensor_flat - reconstructed_flat).view(shape)\n\n        # Store scaling parameters as float32 (v_min, scale)\n        scaling = torch.tensor([v_min.item(), scale.item()], device=tensor.device, dtype=torch.float32)\n\n        # Compressed tensors: quantized values, indices, scaling\n        return [quantized, indices, scaling], shape\n\n    def decompress(self, compressed_tensors, ctx):\n        shape = ctx\n        quantized, indices, scaling = compressed_tensors\n        v_min, scale = scaling[0].item(), scaling[1].item()\n\n        # Dequantize\n        values = quantized.float() * scale + v_min\n\n        # Reconstruct tensor\n        numel = torch.prod(torch.tensor(shape)).item()\n        reconstructed = torch.zeros(numel, dtype=torch.float32, device=values.device)\n        reconstructed.scatter_(0, indices.to(torch.int64), values)\n        return reconstructed.view(shape)"
7}

1{
2  "op": "replace",
3  "filename": "pytorch-vision/custom_compressor.py",
4  "start_line": 182,
5  "end_line": 232,
6  "content": "class Compressor:\n    \"\"\"Adaptive Sparse-Quantized Gradient Compression with Error Feedback.\n\n    Combines TopK sparsification with 8‑bit quantization of the retained values,\n    and error feedback to compensate for the bias introduced by both operations.\n    The compression ratio is interpreted as the fraction of gradient elements\n    to keep (same as TopK). Because quantization reduces the per‑element bit\n    width from 32 to 8, we can retain more elements while staying within the\n    same communication budget. Specifically, we keep\n        k = max(1, int(numel * compress_ratio * extra_factor))\n    where extra_factor = 1.33 (derived from 32→8 bit reduction and index overhead).\n    The retained values are linearly quantized to 8‑bit integers using a per‑tensor\n    scaling factor (max absolute value). Error feedback accumulates the difference\n    between the original gradient (plus previous residual) and the dequantized\n    sparse approximation.\n    \"\"\"\n\n    def __init__(self, compress_ratio=0.01):\n        self.compress_ratio = compress_ratio\n        self.residuals = {}\n        # Factor by which we increase the number of kept elements thanks to 8‑bit quantization.\n        # Baseline TopK sends k elements of 32‑bit values + 64‑bit indices.\n        # We send k' elements of 8‑bit values + 64‑bit indices + one scaling float.\n        # Setting extra_factor = 1.33 keeps the total communicated bits roughly equal\n        # to the baseline TopK while retaining more gradient information.\n        self.extra_factor = 1.33\n\n    def compress(self, tensor, name):\n        # Error feedback: add accumulated residual\n        if name in self.residuals:\n            tensor = tensor + self.residuals[name]\n\n        shape = tensor.shape\n        tensor_flat = tensor.flatten()\n        numel = tensor_flat.numel()\n        # Compute number of elements to keep\n        k = max(1, int(numel * self.compress_ratio * self.extra_factor))\n\n        # Select top‑k by magnitude\n        _, indices = torch.topk(tensor_flat.abs(), k, sorted=False)\n        values = tensor_flat[indices]\n\n        # Linear quantization to 8‑bit signed integer\n        max_abs = values.abs().max()\n        if max_abs == 0:\n            # All selected values are zero; send zeros directly\n            quantized = torch.zeros_like(values, dtype=torch.int8)\n            scale = 1.0\n        else:\n            scale = 127.0 / max_abs\n            quantized = torch.round(values * scale).clamp(-127, 127).to(torch.int8)\n\n        # Update residual: store the difference between the original (with residual)\n        # and the dequantized sparse approximation\n        dequantized = quantized.float() / scale\n        decompressed_flat = torch.zeros_like(tensor_flat)\n        decompressed_flat.scatter_(0, indices, dequantized)\n        self.residuals[name] = tensor_flat - decompressed_flat\n        self.residuals[name] = self.residuals[name].view(shape)\n\n        # compressed_tensors: quantized values (int8), indices (int64), scaling factor (float32)\n        return [quantized, indices, torch.tensor([scale], dtype=torch.float32, device=tensor.device)], (numel, shape)\n\n    def decompress(self, compressed_tensors, ctx):\n        quantized, indices, scale_tensor = compressed_tensors\n        scale = scale_tensor.item()\n        numel, shape = ctx\n\n        # Dequantize\n        values = quantized.float() / scale\n\n        # Reconstruct sparse tensor\n        tensor_decompressed = torch.zeros(numel, dtype=values.dtype, device=values.device)\n        tensor_decompressed.scatter_(0, indices, values)\n        return tensor_decompressed.view(shape)"
7}

1{
2  "op": "replace",
3  "filename": "pytorch-vision/custom_compressor.py",
4  "start_line": 224,
5  "end_line": 243,
6  "content": "          # Stochastic rounding to 8‑bit signed integer (unbiased)\n          max_abs = values.abs().max()\n          if max_abs == 0:\n              # All selected values are zero; send zeros directly\n              quantized = torch.zeros_like(values, dtype=torch.int8)\n              scale = 1.0\n          else:\n              scale = 127.0 / max_abs\n              # Map values to [-127, 127] range\n              scaled = values * scale\n              # Stochastic rounding\n              floor = scaled.floor()\n              prob = torch.rand_like(scaled)\n              is_next = (prob < (scaled - floor)).float()\n              rounded = floor + is_next\n              # Clamp to valid range\n              quantized = rounded.clamp(-127, 127).to(torch.int8)\n\n          # Update residual: store the difference between the original (with residual)\n          # and the dequantized sparse approximation\n          dequantized = quantized.float() / scale\n          decompressed_flat = torch.zeros_like(tensor_flat)\n          decompressed_flat.scatter_(0, indices, dequantized)\n          self.residuals[name] = tensor_flat - decompressed_flat\n          self.residuals[name] = self.residuals[name].view(shape)"
7}

1{
2  "op": "replace",
3  "filename": "pytorch-vision/custom_compressor.py",
4  "start_line": 217,
5  "end_line": 248,
6  "content": "          # Compute number of elements to keep\n          k = max(1, int(numel * self.compress_ratio * self.extra_factor))\n\n          # Select top‑k by magnitude\n          _, indices = torch.topk(tensor_flat.abs(), k, sorted=False)\n          values = tensor_flat[indices]\n\n          # Stochastic rounding to 8‑bit signed integer (unbiased)\n          max_abs = values.abs().max()\n          if max_abs == 0:\n              # All selected values are zero; send zeros directly\n              quantized = torch.zeros_like(values, dtype=torch.int8)\n              scale = 1.0\n          else:\n              scale = 127.0 / max_abs\n              # Map values to [-127, 127] range\n              scaled = values * scale\n              # Stochastic rounding\n              floor = scaled.floor()\n              prob = torch.rand_like(scaled)\n              is_next = (prob < (scaled - floor)).float()\n              rounded = floor + is_next\n              # Clamp to valid range\n              quantized = rounded.clamp(-127, 127).to(torch.int8)\n\n          # Update residual: store the difference between the original (with residual)\n          # and the dequantized sparse approximation\n          dequantized = quantized.float() / scale\n          decompressed_flat = torch.zeros_like(tensor_flat)\n          decompressed_flat.scatter_(0, indices, dequantized)\n          self.residuals[name] = tensor_flat - decompressed_flat\n          self.residuals[name] = self.residuals[name].view(shape)"
7}