Fix Reshape issue when shape size is -1 (#10356)

* Fix Reshape issue (in_place) when shape size is -1
This commit is contained in:
pallavides 2022-01-24 19:30:52 -08:00 committed by GitHub
parent 4b87d2c172
commit 790c3be7e9
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 14 additions and 6 deletions

View file

@ -4,14 +4,14 @@
#include "ort_ops.h"
#include "ort_util.h"
#include "ort_log.h"
#include "core/providers/cpu/tensor/reshape_helper.h"
namespace torch_ort {
namespace eager {
void copy(onnxruntime::ORTInvoker& invoker,
void copy(onnxruntime::ORTInvoker& invoker,
const OrtValue& src, OrtValue& dst){
auto& ort_ep = invoker.GetCurrentExecutionProvider();
const auto& src_tensor = src.Get<onnxruntime::Tensor>();
auto* dst_tensor = dst.GetMutable<onnxruntime::Tensor>();
if (!dst_tensor)
@ -25,16 +25,18 @@ void createInplaceOutputValue(OrtValue& input, V<int64_t> shape, OrtValue* p_mlv
// the ort TensorShape class only accept std::vector, so have to conversion.
std::vector<int64_t> new_shape;
new_shape.assign(shape.begin(), shape.end());
onnxruntime::ReshapeHelper helper(input.Get<onnxruntime::Tensor>().Shape(), new_shape);
CreateMLValue(input_ort_tensor->MutableDataRaw(),
input_ort_tensor->DataType(), new_shape, p_mlvalue);
}
template <typename T>
template <typename T>
using Vector = std::vector<T, std::allocator<T>>;
template <>
void createInplaceOutputValue<Vector>(OrtValue& input, Vector<int64_t> shape, OrtValue* p_mlvalue){
auto* input_ort_tensor = input.GetMutable<onnxruntime::Tensor>();
onnxruntime::ReshapeHelper helper(input.Get<onnxruntime::Tensor>().Shape(), shape);
CreateMLValue(input_ort_tensor->MutableDataRaw(),
input_ort_tensor->DataType(), shape, p_mlvalue);
}
@ -42,4 +44,4 @@ void createInplaceOutputValue<Vector>(OrtValue& input, Vector<int64_t> shape, Or
template void createInplaceOutputValue<c10::ArrayRef>(OrtValue& input, c10::ArrayRef<int64_t> shape, OrtValue* p_mlvalue);
} // namespace eager
} // namespace torch_ort
} // namespace torch_ort

View file

@ -19,19 +19,25 @@ class OrtTensorTests(unittest.TestCase):
ort_ones = cpu_ones.to('ort')
assert ort_ones.is_ort
assert torch.allclose(cpu_ones, ort_ones.cpu())
def test_reshape(self):
cpu_ones = torch.ones(10, 10)
ort_ones = cpu_ones.to('ort')
y = ort_ones.reshape(-1)
assert len(y.size()) == 1
assert y.size()[0] == 100
def test_view(self):
cpu_ones = torch.ones(2048)
ort_ones = cpu_ones.to('ort')
y = ort_ones.view(4, 512)
assert y.size() == (4, 512)
def test_view_neg1(self):
cpu_ones = torch.ones(784, 256)
ort_ones = cpu_ones.to('ort')
y = ort_ones.view(-1)
assert y.size()[0] == 200704
if __name__ == '__main__':
unittest.main()