code: def test_tf(self, use_cpu_only, backend, rank, k):
      # Set TopK to custom TF function
      _set_tf_op("TopKV2", custom_tf_topk)
      shape = np.random.randint(low=3, high=6, size=rank)
      with tf.Graph().as_default() as graph:
          x = tf.placeholder(tf.float32, shape=shape)
          ref = tf.math.top_k(x, k=k, sorted=True)
          ref = (ref[1], ref[0])
          spec, _, _, _ = tf_graph_to_proto(
              graph,
              {x: random_gen(shape, rand_min=-100, rand_max=100)},
              ref,
              backend=backend,
          )
          layers = spec.neuralNetwork.layers
          assert (
              layers[-1].custom is not None
          ), "Expecting a custom layer"
          assert (
              "TopK" == layers[-1].custom.className
          ), "Custom Layer class name mis-match"
          assert (
              k == layers[-1].custom.parameters["k"].intValue
          ), "Incorrect parameter value k"
          assert (
              True == layers[-1].custom.parameters["sorted"].boolValue
          ), "Incorrect parameter value for Sorted"
      # Set TopK to default conversion function
      _set_tf_op("TopKV2", default_tf_topk)
  
file_path: turicreate/src/external/coremltools_wrap/coremltools/coremltools/converters/mil/frontend/tensorflow/test/test_custom_ops.py
imports: from coremltools.converters.mil import testing_reqs
  from coremltools.converters.mil.testing_reqs import *
  from coremltools.converters.mil.frontend.tensorflow.test.testing_utils import make_tf_graph, tf_graph_to_proto, run_compare_tf
  from coremltools.converters.mil.frontend.tensorflow.tf_op_registry import register_tf_op
  from coremltools.converters.mil.frontend.tensorflow.tf_op_registry import _TF_OPS_REGISTRY
  from coremltools.converters.mil.mil.ops.defs._op_reqs import *
  from coremltools.converters.mil.mil import Builder
  from coremltools.converters.mil.mil.types.symbolic import is_symbolic
label_desc: Use image node_id 1 for calling the TestCustomTopK obj's underlying member method code with example usage: obj.test_tf(use_cpu_only, backend, rank, k) without return types