Add pad_variable_length_inputs to BatchOptins and export configs

export/orbax/export/obm_configs.py

@@ -196,6 +196,9 @@ class BatchOptions:
       all batch sizes no larger than `max_batch_size` are allowed. Otherwise,
       supplies a list of batch sizes. The entries must increase monotonically.
     disable_large_batch_splitting: Whether to disable large batch splitting.
+    pad_variable_length_inputs: Whether to dynamically pad variable-length
+      non-batch dimensions to the maximum size across requests within each
+      batch. Requires disable_large_batch_splitting = true.
     batch_padding_policy: The batch padding policy for the batch scheduler.
       Default is PAD_UP.
     low_priority_batch_options: The batch options for low priority inputs.
@@ -218,6 +221,7 @@ class BatchOptions:
   num_batch_threads: int = 1
   max_enqueued_batches: int = 250
   disable_large_batch_splitting: bool = False
+  pad_variable_length_inputs: bool = False
   batch_padding_policy: BatchPaddingPolicy = BatchPaddingPolicy.PAD_UP
   low_priority_batch_options: LowPriorityBatchOptions | None = None
   mixed_priority_batching_policy: MixedPriorityBatchingPolicy = (
@@ -317,6 +321,15 @@ def __post_init__(self):
     if self.batch_component == BatchComponent.NO_BATCHING:
       return
 
+    if (
+        self.pad_variable_length_inputs
+        and not self.disable_large_batch_splitting
+    ):
+      raise ValueError(
+          "disable_large_batch_splitting must be True when"
+          " pad_variable_length_inputs is True."
+      )
+
     if self.max_batch_size is None:
       if self.allowed_batch_sizes:
         self.max_batch_size = self.allowed_batch_sizes[-1]

export/orbax/export/obm_configs_test.py

@@ -26,6 +26,19 @@ def test_batch_options_disable_batch(self):
     )
     self.assertIsNone(batch_options.max_batch_size)
 
+  def test_pad_variable_length_inputs_with_large_batch_splitting_fails(self):
+    with self.assertRaisesRegex(
+        ValueError,
+        "disable_large_batch_splitting must be True when"
+        " pad_variable_length_inputs is True.",
+    ):
+      obm_configs.BatchOptions(
+          batch_component=obm_configs.BatchComponent.MODEL_FUNCTION,
+          max_batch_size=8,
+          pad_variable_length_inputs=True,
+          disable_large_batch_splitting=False,
+      )
+
   def test_batch_options_raise_error_without_max_batch_size_and_allowed_batch_sizes(
       self,
   ):