code: def testUseStepContainerInFilter(self, apply_filter):
      input_data = np.array([[1, 2, 3], [4, 5, 6]], dtype=np.int64)
  
      # Define a predicate that returns true for the first element of
      # the sequence and not the second, and uses `tf.map_fn()`.
      def _predicate(xs):
          squared_xs = map_fn.map_fn(lambda x: x * x, xs)
          summed = math_ops.reduce_sum(squared_xs)
          return math_ops.equal(summed, 1 + 4 + 9)
  
      dataset = dataset_ops.Dataset.from_tensor_slices(
          [[1, 2, 3], [4, 5, 6]]
      )
      dataset = apply_filter(dataset, _predicate)
      self.assertDatasetProduces(
          dataset, expected_output=[input_data[0]]
      )
  
imports: from absl.testing import parameterized
  import numpy
  from tensorflow.python.data.kernel_tests import checkpoint_test_base
  from tensorflow.python.data.kernel_tests import test_base
  from tensorflow.python.data.ops import dataset_ops
  from tensorflow.python.data.ops import options
  from tensorflow.python.framework import combinations
  from tensorflow.python.framework import errors
  from tensorflow.python.framework import sparse_tensor
  from tensorflow.python.ops import map_fn
  from tensorflow.python.ops import math_ops
  from tensorflow.python.platform import test
label_desc: Use image node_id 4 for calling the FilterTest obj's underlying member method code with example usage: obj.testUseStepContainerInFilter(apply_filter) and returns: math_ops