Refactor Reed-Muller decoder for improved tensor handling and error checking; update tests to reflect parameter name changes in ReedMullerCodeEncoder.

Author: selimfirat

kaira/models/fec/decoders/reed_muller_decoder.py

@@ -241,17 +241,19 @@ def decode_block(r_block):
             errors = torch.zeros_like(r_block) if return_errors else None
 
             for i in range(batch_size):
-                # Get the current received word
-                r = r_block[i]
+                # Get the current received word - ensure it's a 1D tensor
+                if r_block.dim() == 3:  # Handle the case when r_block has shape [batch, 1, code_length]
+                    r = r_block[i, 0, :]
+                else:  # Handle the case when r_block has shape [batch, code_length]
+                    r = r_block[i, :]
 
+                """
                 # Convert to binary for hard decoding or compute hard decisions for soft decoding
                 if self.input_type == "hard":
                     bx = r.clone()
                 else:  # self.input_type == "soft"
                     bx = (r < 0).to(torch.int)
-
-                # Original received bits (for error calculation)
-                # original_bits = bx.clone() if return_errors else None
+                """
 
                 # Decode using Reed algorithm
                 u_hat = torch.zeros(self.code_dimension, dtype=torch.int, device=received.device)
@@ -266,10 +268,20 @@ def decode_block(r_block):
                         # Calculate checksums for each group in the partition
                         checksums = []
                         for group in partition:
+                            # Ensure the group indices are valid
+                            valid_indices = group[group < r.shape[0]]
+                            if len(valid_indices) == 0:
+                                continue
+
                             # Take relevant positions and compute parity
-                            group_bits = bx[group]
+                            # Use indexing to select elements from the 1D tensor
+                            group_bits = r[valid_indices].to(torch.int)
                             checksum = torch.sum(group_bits) % 2
-                            checksums.append(checksum)
+                            checksums.append(checksum.item())  # Use .item() to convert tensor to scalar
+
+                        # Skip if no valid checksums
+                        if not checksums:
+                            continue
 
                         # Convert to tensor
                         checksums = torch.tensor(checksums, device=received.device)
@@ -284,17 +296,26 @@ def decode_block(r_block):
                         min_reliabilities = []
 
                         for group in partition:
+                            # Ensure the group indices are valid
+                            valid_indices = group[group < r.shape[0]]
+                            if len(valid_indices) == 0:
+                                continue
+
                             # Take relevant positions
-                            group_bits = bx[group]
-                            group_reliabilities = torch.abs(r[group])
+                            group_bits = (r[valid_indices] < 0).to(torch.int)
+                            group_reliabilities = torch.abs(r[valid_indices])
 
                             # Compute parity of hard decisions
                             checksum = torch.sum(group_bits) % 2
-                            checksums.append(checksum)
+                            checksums.append(checksum.item())  # Use .item() to convert tensor to scalar
 
                             # Find minimum reliability in this group
                             min_reliability = torch.min(group_reliabilities)
-                            min_reliabilities.append(min_reliability)
+                            min_reliabilities.append(min_reliability.item())  # Use .item() to convert tensor to scalar
+
+                        # Skip if no valid checksums
+                        if not checksums:
+                            continue
 
                         # Convert to tensors
                         checksums = torch.tensor(checksums, device=received.device)
@@ -306,18 +327,14 @@ def decode_block(r_block):
                         # Make decision
                         u_hat[j] = (decision_var < 0).to(torch.int)
 
-                    # Cancel the effect of this bit from the received word
-                    # In a complete implementation, this would use the generator matrix
-                    # bx ^= u_hat[j] * self.encoder.generator_matrix[j]
-
                 # Store the decoded message
                 decoded[i] = u_hat
 
                 # Compute error pattern if needed
                 if return_errors:
                     # Re-encode the message to get the correct codeword
-                    correct_codeword = self.encoder(u_hat.unsqueeze(0)).squeeze(0)
-                    errors[i] = (r.to(torch.int) != correct_codeword).to(torch.int)
+                    correct_codeword = self.encoder(u_hat.float().unsqueeze(0)).squeeze(0)
+                    errors[i] = (r.to(torch.int) != correct_codeword.to(torch.int)).to(torch.int)
 
             return (decoded, errors) if return_errors else decoded
 

kaira/models/fec/utils.py

@@ -131,7 +131,7 @@ def apply_blockwise(x: torch.Tensor, block_size: int, fn: Callable) -> torch.Ten
             a transformed tensor preserving the batch dimensions
 
     Returns:
-        Tensor with transformed blocks
+        Tensor with transformed blocks or tuple of tensors if fn returns a tuple
 
     Raises:
         AssertionError: If the last dimension is not divisible by block_size
@@ -152,5 +152,14 @@ def apply_blockwise(x: torch.Tensor, block_size: int, fn: Callable) -> torch.Ten
     # Apply function along the last dimension (block)
     result = fn(x_reshaped)
 
-    # Flatten the result back to original structure
-    return result.view(*leading_dims, -1)
+    # Check if the result is a tuple (like when return_errors=True)
+    if isinstance(result, tuple):
+        # Process each part of the tuple independently
+        processed_results = []
+        for res_part in result:
+            # Flatten each part back to original structure
+            processed_results.append(res_part.view(*leading_dims, -1))
+        return tuple(processed_results)
+    else:
+        # Flatten the result back to original structure
+        return result.view(*leading_dims, -1)

tests/models/fec/decoders/test_reed_muller_decoder.py

@@ -13,7 +13,7 @@ class TestReedMullerDecoder:
     def test_initialization(self):
         """Test initialization with valid parameters."""
         # Create a Reed-Muller encoder for RM(1,3)
-        encoder = ReedMullerCodeEncoder(r=1, m=3)
+        encoder = ReedMullerCodeEncoder(order=1, length_param=3)
 
         # Initialize the decoder with default parameters (hard-decision)
         decoder = ReedMullerDecoder(encoder=encoder)
@@ -35,7 +35,7 @@ def test_initialization(self):
     def test_generate_reed_partitions(self):
         """Test generation of Reed partitions."""
         # Create a Reed-Muller encoder for RM(1,3)
-        encoder = ReedMullerCodeEncoder(r=1, m=3)
+        encoder = ReedMullerCodeEncoder(order=1, length_param=3)
         decoder = ReedMullerDecoder(encoder=encoder)
 
         # Generate Reed partitions
@@ -51,7 +51,7 @@ def test_generate_reed_partitions(self):
     def test_decoding_no_errors_hard_decision(self):
         """Test hard-decision decoding with no errors."""
         # Create a Reed-Muller encoder for RM(1,3)
-        encoder = ReedMullerCodeEncoder(r=1, m=3)
+        encoder = ReedMullerCodeEncoder(order=1, length_param=3)
         decoder = ReedMullerDecoder(encoder=encoder, input_type="hard")
 
         # Create a message and encode it
@@ -71,7 +71,7 @@ def test_decoding_no_errors_hard_decision(self):
     def test_decoding_with_errors_hard_decision(self):
         """Test hard-decision decoding with errors."""
         # Create a Reed-Muller encoder for RM(1,3)
-        encoder = ReedMullerCodeEncoder(r=1, m=3)
+        encoder = ReedMullerCodeEncoder(order=1, length_param=3)
         decoder = ReedMullerDecoder(encoder=encoder, input_type="hard")
 
         # Create a message and encode it
@@ -91,7 +91,7 @@ def test_decoding_with_errors_hard_decision(self):
     def test_decoding_soft_decision(self):
         """Test soft-decision decoding."""
         # Create a Reed-Muller encoder for RM(1,3)
-        encoder = ReedMullerCodeEncoder(r=1, m=3)
+        encoder = ReedMullerCodeEncoder(order=1, length_param=3)
         decoder = ReedMullerDecoder(encoder=encoder, input_type="soft")
 
         # Create a message and encode it
@@ -117,7 +117,7 @@ def test_decoding_soft_decision(self):
     def test_decoding_with_return_errors(self):
         """Test decoding with return_errors=True."""
         # Create a Reed-Muller encoder for RM(1,3)
-        encoder = ReedMullerCodeEncoder(r=1, m=3)
+        encoder = ReedMullerCodeEncoder(order=1, length_param=3)
         decoder = ReedMullerDecoder(encoder=encoder)
 
         # Create a message and encode it
@@ -139,7 +139,7 @@ def test_decoding_with_return_errors(self):
     def test_decoding_with_batch_dimension(self):
         """Test decoding with batch dimension."""
         # Create a Reed-Muller encoder for RM(1,3)
-        encoder = ReedMullerCodeEncoder(r=1, m=3)
+        encoder = ReedMullerCodeEncoder(order=1, length_param=3)
         decoder = ReedMullerDecoder(encoder=encoder)
 
         # Create messages and encode them
@@ -166,7 +166,7 @@ def test_decoding_with_batch_dimension(self):
     def test_invalid_input_dimensions(self):
         """Test decoding with invalid input dimensions."""
         # Create a Reed-Muller encoder for RM(1,3)
-        encoder = ReedMullerCodeEncoder(r=1, m=3)
+        encoder = ReedMullerCodeEncoder(order=1, length_param=3)
         decoder = ReedMullerDecoder(encoder=encoder)
 
         # Create a received word with invalid length
@@ -180,15 +180,15 @@ def test_invalid_input_dimensions(self):
     def test_multiple_rm_parameters(self):
         """Test with different Reed-Muller code parameters."""
         # Test with RM(0,3) - repetition code
-        encoder_rm03 = ReedMullerCodeEncoder(r=0, m=3)
+        encoder_rm03 = ReedMullerCodeEncoder(order=0, length_param=3)
         decoder_rm03 = ReedMullerDecoder(encoder=encoder_rm03)
 
         # For RM(0,3), dimension = 1, length = 8
         assert encoder_rm03.code_dimension == 1
         assert encoder_rm03.code_length == 8
 
         # Test with RM(1,4) - first-order code
-        encoder_rm14 = ReedMullerCodeEncoder(r=1, m=4)
+        encoder_rm14 = ReedMullerCodeEncoder(order=1, length_param=4)
         decoder_rm14 = ReedMullerDecoder(encoder=encoder_rm14)
 
         # For RM(1,4), dimension = 5, length = 16