code: def top_k_ref(self, X, k, flatten_indices, axis=-1):
      in_dims = X.shape
      out_dims = list(in_dims)
      out_dims[axis] = k
      out_dims = tuple(out_dims)
      if axis == -1:
          axis = len(in_dims) - 1
      prev_dims = 1
      next_dims = 1
      for i in range(axis):
          prev_dims *= in_dims[i]
      for i in range(axis + 1, len(in_dims)):
          next_dims *= in_dims[i]
      n = in_dims[axis]
      X_flat = X.reshape((prev_dims, n, next_dims))
  
      values_ref = np.ndarray(
          shape=(prev_dims, k, next_dims), dtype=np.float32
      )
      values_ref.fill(0)
      indices_ref = np.ndarray(
          shape=(prev_dims, k, next_dims), dtype=np.int64
      )
      indices_ref.fill(-1)
      flatten_indices_ref = np.ndarray(
          shape=(prev_dims, k, next_dims), dtype=np.int64
      )
      flatten_indices_ref.fill(-1)
      for i in range(prev_dims):
          for j in range(next_dims):
              kv = []
              for x in range(n):
                  val = X_flat[i, x, j]
                  y = x * next_dims + i * in_dims[axis] * next_dims + j
                  kv.append((val, x, y))
              cnt = 0
              for val, x, y in sorted(
                  kv, key=lambda x: (x[0], -x[1]), reverse=True
              ):
                  values_ref[i, cnt, j] = val
                  indices_ref[i, cnt, j] = x
                  flatten_indices_ref[i, cnt, j] = y
                  cnt += 1
                  if cnt >= k or cnt >= n:
                      break
  
      values_ref = values_ref.reshape(out_dims)
      indices_ref = indices_ref.reshape(out_dims)
      flatten_indices_ref = flatten_indices_ref.flatten()
  
      if flatten_indices:
          return (values_ref, indices_ref, flatten_indices_ref)
      else:
          return (values_ref, indices_ref)
  
file_path: pytorch/caffe2/python/operator_test/top_k_test.py
imports: import hypothesis.strategies
  import numpy
  from caffe2.python import core
  from hypothesis import given, settings
  import caffe2.python.hypothesis_test_util
  import caffe2.python.serialized_test.serialized_test_util
label_desc: Use image node_id 1 for calling the TestTopK obj's underlying member method code with example usage: obj.top_k_ref(X, k, flatten_indices, axis) and returns: values_ref,  indices_ref,  flatten_indices_ref, values_ref,  indices_ref