code: def topk(a, k, axis, keepdims):
      """Chunk and combine function of topk
  
      Extract the k largest elements from a on the given axis.
      If k is negative, extract the -k smallest elements instead.
      Note that, unlike in the parent function, the returned elements
      are not sorted internally.
      """
      assert keepdims is True
      axis = axis[0]
      if abs(k) >= a.shape[axis]:
          return a
  
      a = np.partition(a, -k, axis=axis)
      k_slice = slice(-k, None) if k > 0 else slice(-k)
      return a[
          tuple(
              k_slice if i == axis else slice(None)
              for i in range(a.ndim)
          )
      ]
  
imports: from __future__ import annotations
  import contextlib
  from collections.abc import Container, Iterable, Sequence
  from functools import wraps
  from numbers import Integral
  import numpy
  from tlz import concat
  from dask.core import flatten
label_desc: Use image node_id 4 for calling a global function with example usage: topk(a, k, axis, keepdims) and returns: a, a