code: def testGenerateMask(
      self, segments, max_seq_length, expected, axis=-1, descr=None
  ):
      max_seq_length = constant_op.constant(max_seq_length)
      segments = [ragged_factory_ops.constant(i) for i in segments]
      expected = [ragged_factory_ops.constant(i) for i in expected]
      trimmer = trimmer_ops.WaterfallTrimmer(max_seq_length, axis=axis)
      actual = trimmer.generate_mask(segments)
      for expected_mask, actual_mask in zip(expected, actual):
          self.assertAllEqual(actual_mask, expected_mask)
  
imports: from absl.testing import parameterized
  import numpy
  import tensorflow
  import tensorflow_text
  from tensorflow.lite.python import interpreter
  from tensorflow.python.framework import constant_op
  from tensorflow.python.framework import test_util
  from tensorflow.python.ops.ragged import ragged_factory_ops
  from tensorflow.python.platform import test
  from tensorflow_text.python.ops import trimmer_ops
label_desc: Use image node_id 1 for calling the WaterfallTrimmerOpsTest obj's underlying member method code with example usage: obj.testGenerateMask(segments, max_seq_length, expected, axis, descr) without return types