code: def create_custom_TopK(self):
      # Defining SSA TopK Op
      @register_op(is_custom_op=True)
      class custom_topk(Operation):
          input_spec = InputSpec(
              x=TensorInputType(type_domain="T"),
              k=TensorInputType(
                  const=True, optional=True, type_domain=types.int32
              ),
              axis=TensorInputType(
                  const=True, optional=True, type_domain=types.int32
              ),
              sorted=TensorInputType(
                  const=True, optional=True, type_domain=types.bool
              ),
          )
  
          type_domains = {
              "T": (types.fp16, types.fp32),
          }
  
          bindings = {
              "class_name": "TopK",
              "input_order": ["x"],
              "parameters": ["k", "axis", "sorted"],
              "description": "Top K Custom layer",
          }
  
          def default_inputs(self):
              return DefaultInputs(
                  k=1,
                  axis=-1,
                  sorted=False,
              )
  
          def __init__(self, **kwargs):
              super(custom_topk, self).__init__(**kwargs)
  
          def type_inference(self):
              x_type = self.x.dtype
              x_shape = self.x.shape
              k = self.k.val
              axis = self.axis.val
  
              if not is_symbolic(x_shape[axis]) and k > x_shape[axis]:
                  msg = "K={} is greater than size of the given axis={}"
                  raise ValueError(msg.format(k, axis))
  
              ret_shape = list(x_shape)
              ret_shape[axis] = k
              return types.tensor(x_type, ret_shape), types.tensor(
                  types.int32, ret_shape
              )
  
      # Following logging is to ensure testing of TopK implemented in tf converter
      # default path is testing with appropriate conversion function
      # Log default tf topk
      default_tf_topk = _TF_OPS_REGISTRY.get("TopKV2", None)
  
      # Override TopK op with override=True flag
      @register_tf_op(tf_alias=["TopKV2"], override=True)
      def CustomTopK(context, node):
          x = context[node.inputs[0]]
          k = context[node.inputs[1]]
          sorted = node.attr.get("sorted", False)
          x = mb.custom_topk(
              x=x, k=k.val, axis=-1, sorted=sorted, name=node.name
          )
          context.add(node.name, x)
  
      yield
  
      _TF_OPS_REGISTRY["TopKV2"] = default_tf_topk
  
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 1 for calling the TestCustomTopK obj's underlying member method code with example usage: obj.create_custom_TopK() and returns: DefaultInputs, types,  types