Tensor from Python appear to be copied

[jamesETsmith]

Summary

Hello! I'm playing around with TACO and would like to pass some NumPy tensors to C++ and convert them to TACO's Tensor type before working with them in C++, but I'm finding that the modifications I make on the C++ side of things are not reflected in the original NumPy array. I tried to explicitly avoid copying here so I thought any modifications I made in C++ should show up in the original NumPy array. Am I missing something here?

Minimal Example

Here is a minimal example adapted from taco/python_bindings/src/pyTensor.cpp:

#include "pyTensor.hpp"

#include <iostream>
#include <type_traits>

#include "pybind11/numpy.h"
#include "pybind11/operators.h"
#include "pybind11/stl.h"
#include "taco/tensor.h"
#include "taco/type.h"

//
// Adapted from pyTensor.cpp
//

// Create Tensor from NumPy C-ordered array (this should not copy the NumPy data)
template <typename T>
static Tensor<T> fromNpArr(py::buffer_info &array_buffer, Format &fmt) {
  std::vector<ssize_t> buf_shape = array_buffer.shape;
  std::vector<int> shape(buf_shape.begin(), buf_shape.end());
  const ssize_t size = array_buffer.size;

  // Creat row-major dense tensor
  Tensor<T> tensor(shape, fmt);
  TensorStorage &storage = tensor.getStorage();
  void *buf_data = array_buffer.ptr;
  Array::Policy policy = Array::Policy::UserOwns;
  storage.setValues(makeArray(static_cast<T *>(buf_data), size, policy));
  tensor.setStorage(storage);
  return tensor;
}

template <typename T>
static Tensor<T> fromNumpyC(py::array_t<T, py::array::c_style> &array) {
  py::buffer_info array_buffer = array.request();
  const ssize_t dims = array_buffer.ndim;
  Format fmt(std::vector<ModeFormatPack>(dims, dense));
  return fromNpArr<T>(array_buffer, fmt);
}

//
// End functions adapted from pyTensor.cpp
//

// Modify the TACO tensor and (hopefully) the NumPy array as well
static void modify_tensor(py::array_t<double, py::array::c_style> &array) {
  Tensor<double> tensor = fromNumpyC<double>(array);
  std::cout << tensor(0, 0) << std::endl;
  tensor(0, 0) = 11.0;
  std::cout << tensor(0, 0) << std::endl;
}

// Make Python bindings for `modify_tensor`
PYBIND11_MODULE(taco_tests, m) {
  m.def("modify_tensor", &modify_tensor, "Test function");
}

And a minimal Python test script:

import numpy as np
from taco_tests import modify_tensor, my_modify

np.random.seed(20)
a = np.ascontiguousarray(np.random.random((4, 4)), dtype=np.float64)
print(a[0, 0])
modify_tensor(a)
print(a[0, 0])

Which has the following output:

0.5881308010772742
0.588131
11
0.5881308010772742 # <=== Should be 11

Possible Solutions

Maybe some complicated things going on under the hood with Pybind11?

Any ideas/suggestions, would be much appreciated!

[stephenchouca]

I believe the problem is that TACO currently does not support in-place updates of tensor elements, so tensor(0, 0) = 11.0; will actually allocate a new tensor under the hood. If you really need to update tensor elements in-place, then one option would be to get a reference to the underlying array that stores the values and then modify that directly by doing something like:

double* arr = (double*)tensor.getStorage().getValues().getData();
arr[0] = 11;

(Note that this assumes the tensor is dense; if it is stored in a sparse format, then that'd be more difficult since you'd also have to update the index arrays.) In the long run though, we would definitely like to extend TACO to support in-place updates of tensor elements, which would solve problems like this.