mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-30 20:18:08 +00:00
Fix Reshape issue when shape size is -1 (#10356)
* Fix Reshape issue (in_place) when shape size is -1
This commit is contained in:
parent
4b87d2c172
commit
790c3be7e9
2 changed files with 14 additions and 6 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
Loading…
Reference in a new issue