From 6c067e099856d8e44ecff3280fc55dde17b3494c Mon Sep 17 00:00:00 2001
From: zhangbo9674 <82555433+zhangbo9674@users.noreply.github.com>
Date: Fri, 15 Apr 2022 09:45:15 +0800
Subject: [PATCH] support weakref for eager tensor (#41769) (#41797)

---
 paddle/fluid/pybind/eager.cc | 3 +++
 paddle/fluid/pybind/eager.h  | 2 ++
 2 files changed, 5 insertions(+)

diff --git a/paddle/fluid/pybind/eager.cc b/paddle/fluid/pybind/eager.cc
index c600844596..74d15b6c0c 100644
--- a/paddle/fluid/pybind/eager.cc
+++ b/paddle/fluid/pybind/eager.cc
@@ -709,6 +709,8 @@ int TensorInit(PyObject* self, PyObject* args, PyObject* kwargs) {
 }
 
 static void TensorDealloc(TensorObject* self) {
+  if (self->weakrefs != NULL)
+    PyObject_ClearWeakRefs(reinterpret_cast<PyObject*>(self));
   self->tensor.~Tensor();
   Py_TYPE(self)->tp_free(reinterpret_cast<PyObject*>(self));
 }
@@ -739,6 +741,7 @@ void BindEager(pybind11::module* module) {
   type->tp_getset = variable_properties;
   type->tp_init = TensorInit;
   type->tp_new = TensorNew;
+  type->tp_weaklistoffset = offsetof(TensorObject, weakrefs);
   Py_INCREF(&PyBaseObject_Type);
   type->tp_base = reinterpret_cast<PyTypeObject*>(&PyBaseObject_Type);
   type->tp_flags |=
diff --git a/paddle/fluid/pybind/eager.h b/paddle/fluid/pybind/eager.h
index bb55ef62ee..03676a677a 100644
--- a/paddle/fluid/pybind/eager.h
+++ b/paddle/fluid/pybind/eager.h
@@ -22,6 +22,8 @@ namespace pybind {
 
 typedef struct {
   PyObject_HEAD paddle::experimental::Tensor tensor;
+  // Weak references
+  PyObject* weakrefs;
 } TensorObject;
 
 typedef struct {
-- 
GitLab