Disable input forwarding for tensors fed from python or returned from py_func.

alextp · tensorflower-gardener · commit d9b9fa4ffe85 · 2017-03-29T10:53:09.000-07:00
Change: 151591184
diff --git a/tensorflow/c/c_api.cc b/tensorflow/c/c_api.cc
@@ -135,6 +135,9 @@ class TF_ManagedBuffer : public TensorBuffer {
     proto->set_requested_bytes(rb);
     proto->set_allocator_name(tensorflow::cpu_allocator()->Name());
   }
+
+  // Prevents input forwarding from mutating this buffer.
+  bool OwnsMemory() const override { return false; }
 };
 
 void* allocate_tensor(const char* operation, size_t len) {
diff --git a/tensorflow/core/framework/tensor.cc b/tensorflow/core/framework/tensor.cc
@@ -531,7 +531,7 @@ void Tensor::UnsafeCopyFromInternal(const Tensor& other, DataType dtype,
 // one both for the SubBuffer _and_ the underlying TensorBuffer.
 bool Tensor::RefCountIsOne() const {
   return buf_ != nullptr && buf_->RefCountIsOne() &&
-         buf_->root_buffer()->RefCountIsOne();
+         buf_->root_buffer()->RefCountIsOne() && buf_->OwnsMemory();
 }
 
 // The macro CASES() expands to a switch statement conditioned on
diff --git a/tensorflow/core/framework/tensor.h b/tensorflow/core/framework/tensor.h
@@ -502,6 +502,9 @@ class TensorBuffer : public core::RefCounted {
   T* base() const {
     return reinterpret_cast<T*>(data());
   }
+
+  // Whether this TensorBuffer owns the underlying memory.
+  virtual bool OwnsMemory() const { return true; }
 };
 
 template <typename T>
diff --git a/tensorflow/python/kernel_tests/py_func_test.py b/tensorflow/python/kernel_tests/py_func_test.py
@@ -160,6 +160,14 @@ def testCleanup(self):
         _ = script_ops.py_func(lambda x: x + 1, [c], [dtypes.float32])
     self.assertTrue(script_ops._py_funcs.size() < 100)
 
+  def testAlias(self):
+    with self.test_session():
+      np_array = np.array([1.0, 2.0], dtype=np.float32)
+      tf_array = script_ops.py_func(lambda: np_array, [], [dtypes.float32])
+      value = tf_array + constant_op.constant([2.0, 3.0], dtype=dtypes.float32)
+      value.op.run()
+      self.assertAllEqual(np_array, [1.0, 2.0])
+
   def testBadNumpyReturnType(self):
     with self.test_session():
 
diff --git a/tensorflow/python/lib/core/py_func.cc b/tensorflow/python/lib/core/py_func.cc
@@ -254,6 +254,9 @@ class NumpyTensorBuffer : public TensorBuffer {
     return Tensor(dtype, shape, this);
   }
 
+  // Prevents input forwarding from overwriting this buffer.
+  bool OwnsMemory() const override { return false; }
+
  private:
   PyArrayObject* array_;
   size_t len_;

Original file line number	Diff line number	Diff line change
`@@ -135,6 +135,9 @@ class TF_ManagedBuffer : public TensorBuffer {`
`135`	`135`	`proto->set_requested_bytes(rb);`
`136`	`136`	`proto->set_allocator_name(tensorflow::cpu_allocator()->Name());`
`137`	`137`	`}`
	`138`	`+`
	`139`	`+ // Prevents input forwarding from mutating this buffer.`
	`140`	`+ bool OwnsMemory() const override { return false; }`
`138`	`141`	`};`
`139`	`142`
`140`	`143`	`void* allocate_tensor(const char* operation, size_t len) {`
Original file line number	Diff line number	Diff line change
`@@ -531,7 +531,7 @@ void Tensor::UnsafeCopyFromInternal(const Tensor& other, DataType dtype,`
`531`	`531`	`// one both for the SubBuffer _and_ the underlying TensorBuffer.`
`532`	`532`	`bool Tensor::RefCountIsOne() const {`
`533`	`533`	`return buf_ != nullptr && buf_->RefCountIsOne() &&`
`534`		`- buf_->root_buffer()->RefCountIsOne();`
	`534`	`+ buf_->root_buffer()->RefCountIsOne() && buf_->OwnsMemory();`
`535`	`535`	`}`
`536`	`536`
`537`	`537`	`// The macro CASES() expands to a switch statement conditioned on`
Original file line number	Diff line number	Diff line change
`@@ -502,6 +502,9 @@ class TensorBuffer : public core::RefCounted {`
`502`	`502`	`T* base() const {`
`503`	`503`	`return reinterpret_cast<T*>(data());`
`504`	`504`	`}`
	`505`	`+`
	`506`	`+ // Whether this TensorBuffer owns the underlying memory.`
	`507`	`+ virtual bool OwnsMemory() const { return true; }`
`505`	`508`	`};`
`506`	`509`
`507`	`510`	`template <typename T>`
Original file line number	Diff line number	Diff line change
`@@ -254,6 +254,9 @@ class NumpyTensorBuffer : public TensorBuffer {`
`254`	`254`	`return Tensor(dtype, shape, this);`
`255`	`255`	`}`
`256`	`256`
	`257`	`+ // Prevents input forwarding from overwriting this buffer.`
	`258`	`+ bool OwnsMemory() const override { return false; }`
	`259`	`+`
`257`	`260`	`private:`
`258`	`261`	`PyArrayObject* array_;`
`259`	`262`	`size_t len_;`