def METHOD_NAME(fn, opt):
  def wrapper(self, grad, *args, **kwargs):  # pylint: disable=unused-argument
    all_zeros = _is_all_zeros(grad)
    def call_fn():
      with ops.control_dependencies([fn(grad, *args, **kwargs)]):
        return control_flow_ops.no_op()
    return control_flow_ops.cond(all_zeros, control_flow_ops.no_op, call_fn)
  wrapper = types.MethodType(wrapper, opt)
  return wrapper