code: def test_bucket_by_length(self):
      def fake_generator(length, num_examples=1):
          for _ in range(num_examples):
              yield (np.ones((length,)), np.ones((length,)))
  
      def length_function(example):
          return max(example[0].shape[0], example[1].shape[0])
  
      batches = list(
          data.bucket_by_length(
              fake_generator(5, 6),
              length_function,
              [20],
              [2],
              strict_pad_on_len=True,
          )
      )
  
      # We'll get three batches of 2 examples each.
      self.assertLen(batches, 3)
      self.assertIsInstance(batches[0], tuple)
      self.assertLen(batches[0], 2)
      self.assertEqual((2, 20), batches[0][0].shape)
      self.assertEqual((2, 20), batches[0][1].shape)
  
imports: import itertools
  import os
  from absl.testing import absltest
  from absl.testing import parameterized
  import numpy
  from trax import data
label_desc: Use image node_id 31 for calling the InputsTest obj's underlying member method code with example usage: obj.test_bucket_by_length() and returns: max