code: def _with_keepdims(func):
      def _wrapper(data, axis=None, keepdims=False):
          if not keepdims:
              return func(data, axis=axis)
          else:
              if axis is not None:
                  axis = axis if isinstance(axis, int) else axis[0]
                  out_shape = list(data.shape)
                  out_shape[axis] = 1
              else:
                  out_shape = [1 for _ in range(len(data.shape))]
              return func(data, axis=axis).reshape(out_shape)
  
      return _wrapper
  
imports: import sys
  import numpy
  import numpy.random
  import pytest
  import tvm
  import tvm.testing
  import tvm.topi.testing
  from tvm import relay, te
  from tvm.relay import transform
  from tvm.relay.testing import run_infer_type
label_desc: Use image node_id 6 for calling a global function with example usage: _with_keepdims(func) and returns: _wrapper, func, func