Function PyTensorObject_zero_grad

oneflow/api/python/framework/tensor.cpp:340–367

static constexpr auto* SbpBToP = DECORATE(&RawSbpBToP, ThreadLocalCached);

static PyObject* PyTensorObject_zero_grad(PyObject* self, PyObject* args, PyObject* kwargs) {
  HANDLE_ERRORS
  int set_to_none = 0;
  static const char* keywords[2] = {"set_to_none", NULL};
  if (!PyArg_ParseTupleAndKeywords(args, kwargs, "|p:_zero_grad_", const_cast<char**>(keywords),
                                   &set_to_none)) {
    return NULL;
  }
  const auto& t = PyTensor_Unpack(self);
  const auto acc_grad = ASSERT_PTR(t->acc_grad());
  if (acc_grad) {
    if (set_to_none) {
      ASSERT(t->set_acc_grad(NULL));
    } else {
      ASSERT(EagerLocalTensorZeros(acc_grad));
      if (acc_grad->is_global() && acc_grad->is_eager()) {
        const auto local_tensor = ASSERT_PTR(functional::GlobalToLocal(acc_grad, false));
        const auto p = ASSERT_PTR(functional::LocalToGlobal(
            local_tensor, ASSERT(acc_grad->parallel_desc()), SbpBToP(ASSERT(acc_grad->nd_sbp())),
            *acc_grad->shape(), acc_grad->dtype(), false, false));
        ASSERT(acc_grad->set_data(p));
      }
    }
  }
  Py_XINCREF(self);
  return self;
  END_HANDLE_ERRORS
}

static PyObject* PyTensorObject_register_hook(PyObject* self, PyObject* hook) {
  HANDLE_ERRORS