code: def test_masked_input(self, compute_unit, backend):
      input_shape = [4, 10, 8]
      val = np.random.rand(*input_shape).astype(np.float32)
  
      @make_tf_graph([input_shape])
      def build_model(input):
          sliced_input = input[..., 4]
          mask = tf.where(sliced_input > 0)
          masked_input = tf.gather_nd(input, mask)
          return masked_input
  
      model, inputs, outputs = build_model
  
      input_values = [val]
      input_dict = dict(zip(inputs, input_values))
      TensorFlowBaseTest.run_compare_tf(
          model,
          input_dict,
          outputs,
          compute_unit=compute_unit,
          backend=backend,
      )
  
file_path: coremltools/coremltools/converters/mil/frontend/tensorflow/test/test_graphs.py
imports: import itertools
  import numpy
  import pytest
  from coremltools.converters.mil.frontend.tensorflow.test.testing_utils import TensorFlowBaseTest, make_tf_graph
  from coremltools.converters.mil.testing_reqs import backends, compute_units
label_desc: Use image node_id 1 for calling the TestTFGraphs obj's underlying member method code with example usage: obj.test_masked_input(compute_unit, backend) and returns: masked_input