feat(compression): add SpecBuilder for programmatic compression specs (#3133)

Add a fluent builder API for creating compression specifications
without writing YAML strings. This is useful in scripts and
Jupyter notebooks.

Example usage:
    spec = (compression.SpecBuilder()
        .add_tensor(subgraph=0, tensor=2)
            .with_lut(index_bitwidth=4)
        .build())

BUG=#3125

Co-authored-by: suleshahid <110432064+suleshahid@users.noreply.github.com>

Author: rkuester

python/tflite_micro/postinstall_check.py

@@ -61,6 +61,10 @@ def compression_test():
   # with compressible tensors, but we verify the function is importable
   assert callable(compression.compress)
 
+  # Test availability of the SpecBuilder
+  _ = (compression.SpecBuilder().add_tensor(
+      subgraph=0, tensor=0).with_lut(index_bitwidth=4).build())
+
   return True
 
 

tensorflow/lite/micro/compression/BUILD

@@ -23,6 +23,7 @@ py_library(
     deps = [
         ":compress_lib",
         ":spec",
+        ":spec_builder",
     ],
 )
 
@@ -189,6 +190,25 @@ py_test(
     ],
 )
 
+py_library(
+    name = "spec_builder",
+    srcs = ["spec_builder.py"],
+    deps = [
+        ":spec",
+    ],
+)
+
+py_test(
+    name = "spec_builder_test",
+    size = "small",
+    srcs = ["spec_builder_test.py"],
+    deps = [
+        ":spec",
+        ":spec_builder",
+        requirement("tensorflow"),
+    ],
+)
+
 py_library(
     name = "test_models",
     srcs = ["test_models.py"],

tensorflow/lite/micro/compression/__init__.py

@@ -22,5 +22,6 @@
 
 from .compress import compress
 from .spec import parse_yaml
+from .spec_builder import SpecBuilder
 
-__all__ = ["compress", "parse_yaml"]
+__all__ = ["compress", "parse_yaml", "SpecBuilder"]

tensorflow/lite/micro/compression/spec_builder.py

@@ -0,0 +1,101 @@
+# Copyright 2025 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+"""Builder pattern for creating compression specifications programmatically.
+
+This module provides a fluent API for building compression specs without
+needing to write YAML strings.
+
+Example usage:
+    from tflite_micro.compression import SpecBuilder
+    
+    spec = (SpecBuilder()
+        .add_tensor(subgraph=0, tensor=2)
+            .with_lut(index_bitwidth=4)
+        .add_tensor(subgraph=0, tensor=4)
+            .with_lut(index_bitwidth=2)
+        .build())
+"""
+
+from typing import List, Optional
+from . import spec
+
+
+class TensorBuilder:
+  """Builder for individual tensor compression specifications."""
+
+  def __init__(self, subgraph: int, tensor: int,
+               parent_builder: 'SpecBuilder'):
+    self.subgraph = subgraph
+    self.tensor = tensor
+    self.compression_methods: List[spec.CompressionMethod] = []
+    self._parent = parent_builder
+
+  def with_lut(self, index_bitwidth: int) -> 'SpecBuilder':
+    """Add LUT compression to this tensor.
+        
+        Args:
+            index_bitwidth: Number of bits for the LUT index (e.g., 4 for 16 values)
+            
+        Returns:
+            The parent SpecBuilder for method chaining
+        """
+    self.compression_methods.append(
+        spec.LookUpTableCompression(index_bitwidth=index_bitwidth))
+    return self._parent
+
+  def _build(self) -> spec.Tensor:
+    """Build the Tensor specification object."""
+    return spec.Tensor(subgraph=self.subgraph,
+                       tensor=self.tensor,
+                       compression=self.compression_methods)
+
+
+class SpecBuilder:
+  """Fluent builder for compression specifications."""
+
+  def __init__(self):
+    self._tensor_builders: List[TensorBuilder] = []
+    self._current_tensor: Optional[TensorBuilder] = None
+
+  def add_tensor(self, subgraph: int, tensor: int) -> TensorBuilder:
+    """Add a tensor to be compressed.
+        
+        Args:
+            subgraph: The subgraph index containing the tensor
+            tensor: The tensor index within the subgraph
+            
+        Returns:
+            A TensorBuilder for configuring compression methods
+        """
+    # Finalize any current tensor
+    if self._current_tensor is not None:
+      self._tensor_builders.append(self._current_tensor)
+
+    # Create new tensor builder
+    self._current_tensor = TensorBuilder(subgraph, tensor, self)
+    return self._current_tensor
+
+  def build(self) -> List[spec.Tensor]:
+    """Build the final compression specification.
+        
+        Returns:
+            A list of Tensor specifications ready for use with compress()
+        """
+    # Make sure to include the last tensor if there is one
+    if self._current_tensor is not None:
+      self._tensor_builders.append(self._current_tensor)
+      self._current_tensor = None
+
+    return [tb._build() for tb in self._tensor_builders]

tensorflow/lite/micro/compression/spec_builder_test.py

@@ -0,0 +1,112 @@
+# Copyright 2025 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+"""Tests for the compression spec builder."""
+
+import tensorflow as tf
+
+from tflite_micro.tensorflow.lite.micro.compression import spec
+from tflite_micro.tensorflow.lite.micro.compression import spec_builder
+
+
+class SpecBuilderTest(tf.test.TestCase):
+
+  def test_basic_builder_pattern(self):
+    """Test basic fluent builder usage."""
+    result = (spec_builder.SpecBuilder().add_tensor(
+        subgraph=0, tensor=2).with_lut(index_bitwidth=4).add_tensor(
+            subgraph=0, tensor=4).with_lut(index_bitwidth=2).build())
+
+    self.assertEqual(len(result), 2)
+
+    # Check first tensor
+    self.assertEqual(result[0].subgraph, 0)
+    self.assertEqual(result[0].tensor, 2)
+    self.assertEqual(len(result[0].compression), 1)
+    self.assertIsInstance(result[0].compression[0],
+                          spec.LookUpTableCompression)
+    self.assertEqual(result[0].compression[0].index_bitwidth, 4)
+
+    # Check second tensor
+    self.assertEqual(result[1].subgraph, 0)
+    self.assertEqual(result[1].tensor, 4)
+    self.assertEqual(len(result[1].compression), 1)
+    self.assertIsInstance(result[1].compression[0],
+                          spec.LookUpTableCompression)
+    self.assertEqual(result[1].compression[0].index_bitwidth, 2)
+
+  def test_non_chained_usage(self):
+    """Test using builder without method chaining."""
+    builder = spec_builder.SpecBuilder()
+    builder.add_tensor(0, 2).with_lut(4)
+    builder.add_tensor(0, 4).with_lut(2)
+    result = builder.build()
+
+    self.assertEqual(len(result), 2)
+    self.assertEqual(result[0].tensor, 2)
+    self.assertEqual(result[0].compression[0].index_bitwidth, 4)
+    self.assertEqual(result[1].tensor, 4)
+    self.assertEqual(result[1].compression[0].index_bitwidth, 2)
+
+  def test_empty_spec(self):
+    """Test building an empty spec."""
+    result = spec_builder.SpecBuilder().build()
+    self.assertEqual(len(result), 0)
+
+  def test_single_tensor(self):
+    """Test building a spec with just one tensor."""
+    result = (spec_builder.SpecBuilder().add_tensor(
+        subgraph=2, tensor=42).with_lut(index_bitwidth=16).build())
+
+    self.assertEqual(len(result), 1)
+    self.assertEqual(result[0].subgraph, 2)
+    self.assertEqual(result[0].tensor, 42)
+    self.assertEqual(result[0].compression[0].index_bitwidth, 16)
+
+  def test_tensor_without_compression(self):
+    """Test that tensors can be added without compression methods."""
+    builder = spec_builder.SpecBuilder()
+    # Add tensor but don't call with_lut
+    builder.add_tensor(0, 1)
+    builder.add_tensor(0, 2).with_lut(4)
+    result = builder.build()
+
+    self.assertEqual(len(result), 2)
+    self.assertEqual(result[0].tensor, 1)
+    self.assertEqual(len(result[0].compression), 0)
+    self.assertEqual(result[1].tensor, 2)
+    self.assertEqual(len(result[1].compression), 1)
+
+  def test_builder_produces_same_type_as_parse_yaml(self):
+    """Test that builder produces same data structure as parse_yaml."""
+    # Build using the builder
+    built_spec = (spec_builder.SpecBuilder().add_tensor(
+        subgraph=0, tensor=42).with_lut(index_bitwidth=4).add_tensor(
+            subgraph=0, tensor=55).with_lut(index_bitwidth=2).build())
+
+    # Parse the example YAML from spec.py
+    parsed_spec = spec.parse_yaml(spec.EXAMPLE_YAML_SPEC)
+
+    # They should be equivalent
+    self.assertEqual(len(built_spec), len(parsed_spec))
+    for built, parsed in zip(built_spec, parsed_spec):
+      self.assertEqual(built.subgraph, parsed.subgraph)
+      self.assertEqual(built.tensor, parsed.tensor)
+      self.assertEqual(len(built.compression), len(parsed.compression))
+      self.assertEqual(built.compression[0].index_bitwidth,
+                       parsed.compression[0].index_bitwidth)
+
+
+if __name__ == "__main__":
+  tf.test.main()