code: def test_tf(self, compute_unit, backend, rank, k):
      if backend[0] == "mlprogram":
          pytest.skip(
              "Custom layer not supported with ML Program backend"
          )
  
      input_shape = np.random.randint(low=3, high=6, size=rank)
  
      @make_tf_graph([input_shape])
      def build_model(x):
          ref = tf.math.top_k(x, k=k, sorted=True)
          return ref[1], ref[0]
  
      model, inputs, outputs = build_model
      input_values = [random_gen(input_shape, -1.0, 1.0)]
      input_dict = dict(zip(inputs, input_values))
      spec, _, _, _, _, _ = TensorFlowBaseTest.run_compare_tf(
          model,
          input_dict,
          outputs,
          compute_unit=compute_unit,
          frontend_only=True,
          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 mismatch"
      assert (
          k == layers[-1].custom.parameters["k"].intValue
      ), "Incorrect parameter value k"
      assert (
          layers[-1].custom.parameters["sorted"].boolValue is True
      ), "Incorrect parameter value for Sorted"
  
file_path: coremltools/coremltools/converters/mil/frontend/tensorflow/test/test_custom_ops.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.frontend.tensorflow.tf_op_registry import _TF_OPS_REGISTRY, register_tf_op
  from coremltools.converters.mil.mil import Builder
  from coremltools.converters.mil.mil import Operation, types
  from coremltools.converters.mil.mil.input_type import DefaultInputs, InputSpec, TensorInputType
  from coremltools.converters.mil.mil.ops.defs._op_reqs import register_op
  from coremltools.converters.mil.mil.types.symbolic import is_symbolic
  from coremltools.converters.mil.testing_reqs import backends, compute_units
  from coremltools.converters.mil.testing_utils import random_gen
label_desc: Use image node_id 2 for calling the TestCustomTopK obj's underlying member method code with example usage: obj.test_tf(compute_unit, backend, rank, k) and returns: ref,  ref