fix: make PerGroupCodebook reachable and lossless on packed FP4

crates/ptwm-core/src/codecs/per_group_codebook.rs

@@ -405,10 +405,28 @@ impl PlaneCodec for PerGroupCodebook {
         CodecId::PerGroupCodebook
     }
 
+    /// Accept expanded FP4 value planes whose length divides into whole
+    /// groups.
+    ///
+    /// The codec models one nibble per byte, and `GROUP_SIZE` of 32 is one
+    /// MXFP4 block: the unit that shares a scale, and therefore the unit
+    /// worth giving its own codebook. `MxFp4Deinterleave` produces exactly
+    /// that, 32 single-nibble slots per block.
+    ///
+    /// Two exclusions are deliberate. A *packed* plane, two values per
+    /// byte, is declined rather than misread: reading one would take low
+    /// nibbles only and return them with the high half zeroed, silently
+    /// losing half the data. Such a plane must be deinterleaved first. And
+    /// a length that is not a whole number of groups is declined here
+    /// rather than failing inside `encode`, because dispatch drops the
+    /// whole candidate chain when a codec errors and packed FP4 has only
+    /// one chain to drop. Declining keeps an `encode` error meaning what it
+    /// should: an invariant was violated, not merely an unsuitable input.
     fn accepts(&self, descriptor: &PlaneDescriptor) -> bool {
-        descriptor.is_nibble_packed
+        !descriptor.is_nibble_packed
             && matches!(descriptor.role, Role::Value { .. })
             && descriptor.element_width == ElementWidth::Nibble
+            && (descriptor.length_bytes as usize).is_multiple_of(GROUP_SIZE)
     }
 
     fn priority_for(&self, descriptor: &PlaneDescriptor) -> i8 {
@@ -793,27 +811,27 @@ mod capability_tests {
     }
 
     #[test]
-    fn pgc_accepts_nibble_packed_value_nibble() {
+    fn pgc_accepts_expanded_value_nibble() {
         let c = PerGroupCodebook;
         let d = descriptor(
             Role::Value {
                 format: ValueFormat::Fp4E2m1,
             },
             ElementWidth::Nibble,
-            true,
+            false,
         );
         assert!(c.accepts(&d));
     }
 
     #[test]
-    fn pgc_rejects_not_nibble_packed() {
+    fn pgc_rejects_packed_plane() {
         let c = PerGroupCodebook;
         let d = descriptor(
             Role::Value {
                 format: ValueFormat::Fp4E2m1,
             },
             ElementWidth::Nibble,
-            false,
+            true,
         );
         assert!(!c.accepts(&d));
     }
@@ -826,7 +844,7 @@ mod capability_tests {
                 format: ScaleFormat::E4M3,
             },
             ElementWidth::Nibble,
-            true,
+            false,
         );
         assert!(!c.accepts(&d));
     }
@@ -839,7 +857,7 @@ mod capability_tests {
                 format: ValueFormat::Fp4E2m1,
             },
             ElementWidth::Byte,
-            true,
+            false,
         );
         assert!(!c.accepts(&d));
     }
@@ -852,7 +870,7 @@ mod capability_tests {
                 format: ValueFormat::Fp4E2m1,
             },
             ElementWidth::Nibble,
-            true,
+            false,
         );
         assert_eq!(c.priority_for(&d), 10);
     }
@@ -865,8 +883,131 @@ mod capability_tests {
                 format: ScaleFormat::E4M3,
             },
             ElementWidth::Nibble,
-            true,
+            false,
         );
         assert_eq!(c.priority_for(&d), i8::MIN);
     }
 }
+
+/// End-to-end coverage for the producer this codec is actually built for.
+///
+/// The suite previously exercised the codec only on hand-built planes. It
+/// never drove the op that supplies it in practice, which is how a
+/// descriptor mismatch between the two survived: `MxFp4Deinterleave`
+/// labelled its output nibble-packed while emitting one nibble per byte,
+/// and the codec's own contract expects the expanded form.
+#[cfg(test)]
+mod deinterleave_integration_tests {
+    use super::*;
+    use crate::transforms::mxfp4_deinterleave::MxFp4Deinterleave;
+    use crate::transforms::op::{Op, Plane};
+    use crate::types::descriptor::{Layout, PlaneDescriptor};
+    use crate::types::role::Role;
+
+    const BLOCK_BYTES: usize = 17;
+
+    /// OCP MXFP4 wire bytes: 16 packed-nibble bytes then one E8M0 scale.
+    fn mxfp4_blocks(n_blocks: usize, seed: u64) -> Vec<u8> {
+        let mut s = seed.wrapping_mul(0x9E37_79B9_7F4A_7C15) | 1;
+        let mut out = Vec::with_capacity(n_blocks * BLOCK_BYTES);
+        for b in 0..n_blocks {
+            for _ in 0..16 {
+                s ^= s << 13;
+                s ^= s >> 7;
+                s ^= s << 17;
+                // Skewed within a block, varying across blocks, so a
+                // per-block codebook has structure to find.
+                let lo = (s % 4 + (b as u64 % 3) * 4) as u8 & 0x0F;
+                let hi = ((s >> 8) % 4 + (b as u64 % 3) * 4) as u8 & 0x0F;
+                out.push(lo | (hi << 4));
+            }
+            out.push(0x80 | (b as u8 & 0x0F));
+        }
+        out
+    }
+
+    fn raw_descriptor(len: usize) -> PlaneDescriptor {
+        PlaneDescriptor {
+            role: Role::Raw,
+            element_width: ElementWidth::Byte,
+            length_bytes: len as u64,
+            layout: Layout::Flat,
+            derives_from_tensor: None,
+            residual_of: None,
+            is_nibble_packed: false,
+            vendor_bytes: vec![],
+        }
+    }
+
+    #[test]
+    fn codec_accepts_and_round_trips_the_deinterleaved_value_plane() {
+        let n_blocks = 8;
+        let raw = mxfp4_blocks(n_blocks, 11);
+        let op = MxFp4Deinterleave::new(32).unwrap();
+
+        let descs = op
+            .propagate_descriptors(&[raw_descriptor(raw.len())])
+            .unwrap();
+        let value_desc = descs[0].clone();
+
+        // One group per MXFP4 block is the whole premise of the codec.
+        assert_eq!(value_desc.length_bytes as usize, n_blocks * GROUP_SIZE);
+
+        let c = PerGroupCodebook;
+        assert!(
+            c.accepts(&value_desc),
+            "the codec must accept the plane its own producer emits: {value_desc:?}"
+        );
+
+        let planes = op
+            .forward(&[Plane::new(raw.clone(), raw_descriptor(raw.len())).unwrap()])
+            .unwrap();
+        let value_plane = planes[0].bytes.to_vec();
+        assert_eq!(value_plane.len(), n_blocks * GROUP_SIZE);
+        assert!(
+            value_plane.iter().all(|&b| b <= 0x0F),
+            "the value plane is one nibble per byte, so no byte exceeds 0x0F"
+        );
+
+        let enc = c.encode(&value_plane, None, &PlaneLayout::Flat).unwrap();
+        let dec = c
+            .decode(
+                enc.state_format_version,
+                &enc.state_bytes,
+                &enc.payload,
+                &PlaneLayout::Flat,
+                value_plane.len(),
+            )
+            .unwrap();
+        assert_eq!(
+            dec, value_plane,
+            "codec must be lossless on its own producer"
+        );
+    }
+
+    #[test]
+    fn codec_declines_a_genuinely_packed_plane_instead_of_halving_it() {
+        // Two values per byte. Reading this as one-nibble-per-byte would
+        // drop every high nibble, which is exactly the silent loss the
+        // `accepts` guard exists to prevent.
+        let mut d = raw_descriptor(64);
+        d.role = Role::Value {
+            format: crate::types::role::ValueFormat::Fp4E2m1,
+        };
+        d.element_width = ElementWidth::Nibble;
+        d.is_nibble_packed = true;
+        assert!(!PerGroupCodebook.accepts(&d));
+    }
+
+    #[test]
+    fn codec_declines_a_length_that_is_not_whole_groups() {
+        // Declined rather than failed in `encode`: packed FP4 has one
+        // candidate chain, and a codec error drops the chain entirely.
+        let mut d = raw_descriptor(GROUP_SIZE + 1);
+        d.role = Role::Value {
+            format: crate::types::role::ValueFormat::Fp4E2m1,
+        };
+        d.element_width = ElementWidth::Nibble;
+        assert!(!PerGroupCodebook.accepts(&d));
+    }
+}

crates/ptwm-core/src/compressor.rs

@@ -146,12 +146,17 @@ fn legacy_plane_layout(layout: Layout) -> PlaneLayout {
     }
 }
 
-/// Build a default source `PlaneDescriptor` keyed on dtype code. The
-/// `element_width` follows the canonical mapping used elsewhere in the
-/// crate (see `transforms::source::bytes_per_element` and the byte-/word-
-/// aware transforms). `length_bytes` is the raw byte count;
-/// `Layout::Rows{row_len}` applies when `shape` is supplied, using the
-/// row-major last dimension.
+/// Build a default source `PlaneDescriptor` keyed on dtype code.
+///
+/// `element_width` and `is_nibble_packed` come from `transforms::source`,
+/// which owns the dtype-code mapping. This function used to inline its own
+/// copy of that match. The copy drifted: it kept packed FP4 in the byte
+/// catch-all, so every plane reaching the trial encode was byte-width and
+/// not nibble-packed, and `PerGroupCodebook` could accept none of them.
+/// Call the mapping, do not restate it.
+///
+/// `length_bytes` is the raw byte count; `Layout::Rows{row_len}` applies
+/// when `shape` is supplied, using the row-major last dimension.
 ///
 /// Public so the PyO3 binding (`ptwm-py`) can build the same source
 /// descriptor without duplicating the dtype-code → element-width mapping.
@@ -160,16 +165,7 @@ pub fn source_descriptor_for(
     raw_byte_count: u64,
     shape: Option<&[u64]>,
 ) -> PlaneDescriptor {
-    let element_width = match dtype_code {
-        // FP16 / BF16 / int16 / uint16
-        0x0002 | 0x000F | 0x0007 | 0x0008 => ElementWidth::Word2,
-        // FP32 / int32 / uint32
-        0x0003 | 0x0009 | 0x000A => ElementWidth::Word4,
-        // FP64 / int64 / uint64
-        0x0004 | 0x000B | 0x000C => ElementWidth::Word8,
-        // int8 / uint8 / FP8 variants — and the catch-all
-        _ => ElementWidth::Byte,
-    };
+    let element_width = crate::transforms::source::element_width_for(dtype_code);
     let layout = match shape {
         Some(s) if !s.is_empty() => match u32::try_from(*s.last().unwrap()) {
             Ok(row_len) => Layout::rows(row_len).unwrap_or(Layout::Flat),
@@ -191,7 +187,7 @@ pub fn source_descriptor_for(
         layout,
         derives_from_tensor: None,
         residual_of: None,
-        is_nibble_packed: false,
+        is_nibble_packed: crate::transforms::source::is_nibble_packed_dtype(dtype_code),
         vendor_bytes: vec![],
     }
 }
@@ -1687,3 +1683,48 @@ mod tests {
         }
     }
 }
+
+#[cfg(test)]
+mod source_descriptor_agreement_tests {
+    use super::*;
+    use crate::transforms::op::Op;
+    use crate::transforms::source::Source;
+
+    /// The compress side builds the source descriptor from the raw byte
+    /// count; the decompress side rebuilds it from the chain's Source node
+    /// (shape + dtype code). The two must agree on `length_bytes` for every
+    /// dtype, or a plane decodes at the wrong length.
+    fn assert_paths_agree(dtype_code: u16, shape: &[u64], raw_len: u64) {
+        let from_compressor = source_descriptor_for(dtype_code, raw_len, Some(shape));
+        let src = Source {
+            shape: shape.iter().map(|&d| d as u32).collect(),
+            dtype_code,
+        };
+        let from_chain = &src.propagate_descriptors(&[]).unwrap()[0];
+        assert_eq!(
+            from_compressor.length_bytes, from_chain.length_bytes,
+            "dtype 0x{dtype_code:04X}: compress path says {} bytes, decompress path says {}",
+            from_compressor.length_bytes, from_chain.length_bytes
+        );
+        assert_eq!(from_compressor.element_width, from_chain.element_width);
+        assert_eq!(
+            from_compressor.is_nibble_packed,
+            from_chain.is_nibble_packed
+        );
+    }
+
+    #[test]
+    fn packed_fp4_source_descriptor_agrees_across_both_paths() {
+        // A packed-FP4 tensor's shape counts packed bytes (one byte per
+        // element, two fp4 values), matching `Dtype::element_size` and the
+        // shape a torch `float4_e2m1fn_x2` tensor reports.
+        assert_paths_agree(0x001F, &[128, 128], 128 * 128);
+    }
+
+    #[test]
+    fn byte_and_word_source_descriptors_agree_across_both_paths() {
+        assert_paths_agree(0x0006, &[10, 10], 100);
+        assert_paths_agree(0x0002, &[10, 10], 200);
+        assert_paths_agree(0x0003, &[10, 10], 400);
+    }
+}

crates/ptwm-core/src/dispatch.rs

@@ -125,21 +125,21 @@ mod tests {
     }
 
     #[test]
-    fn dispatch_nibble_value_packed_returns_pgc_first() {
+    fn dispatch_expanded_nibble_value_returns_pgc_first() {
         let d = descriptor(
             Role::Value {
                 format: ValueFormat::Fp4E2m1,
             },
             ElementWidth::Nibble,
             Layout::Flat,
-            true, // is_nibble_packed
+            false, // expanded: one nibble per byte, which is what PGC models
         );
         let result = dispatch(&d);
         assert!(!result.is_empty());
         assert_eq!(
             result[0],
             CodecId::PerGroupCodebook,
-            "PGC must be first for nibble-packed Value + Nibble; got {:?}",
+            "PGC must be first for expanded Value + Nibble; got {:?}",
             result
         );
         assert!(result.contains(&CodecId::Huffman));

crates/ptwm-core/src/transforms/mxfp4_deinterleave.rs

@@ -95,7 +95,15 @@ impl Op for MxFp4Deinterleave {
                 layout: value_layout,
                 derives_from_tensor: inp.derives_from_tensor,
                 residual_of: None,
-                is_nibble_packed: true,
+                // One nibble per byte: `forward` pushes each input byte's
+                // low and high halves as separate slots, so nothing shares
+                // a byte here. The flag means "two values share a byte",
+                // which describes this op's *input*; `ElementWidth::Nibble`
+                // already says these are 4-bit values. Claiming both left
+                // every consumer to choose between reading the plane as
+                // packed, which drops half the data, and reading it as
+                // expanded, which contradicted the flag.
+                is_nibble_packed: false,
                 vendor_bytes: vec![],
             },
             // Plane 1: SCALE — one E8M0 byte per block.
@@ -300,7 +308,8 @@ mod tests {
             }
         );
         assert_eq!(outs[0].element_width, ElementWidth::Nibble);
-        assert!(outs[0].is_nibble_packed);
+        // One nibble per byte: nothing shares a byte in this op's output.
+        assert!(!outs[0].is_nibble_packed);
         assert_eq!(outs[0].length_bytes, 32); // BLOCK_VALUES
     }