code: def ref_top_k_grad(k, abs, base_axis, grad):
      outer_dim = np.prod(grad.shape[:base_axis], dtype=int)
      inner_dim = np.prod(grad.shape[base_axis:], dtype=int)
      gg = grad.reshape(outer_dim, inner_dim).copy()
      ix = np.argsort(np.abs(gg) if abs else gg)[:, -k:]
      dx = np.zeros((outer_dim, inner_dim))
      for idx, row in enumerate(ix):
          dx[idx, row] = gg[idx, row]
      dx = dx.squeeze(axis=0) if base_axis == 0 else dx
      return dx.reshape(grad.shape)
  
file_path: nnabla/python/test/function/test_top_k_grad.py
imports: import pytest
  import numpy
  import nnabla.functions
  from nbla_test_utils import list_context, function_tester
label_desc: Use image node_id 1 for calling a global function with example usage: ref_top_k_grad(k, abs, base_axis, grad) and returns: dx