saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-30 20:18:08 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Add int64_t support for split (#2944)

Browse source 
* Add int64_t support for split
This commit is contained in:
Yufeng Li 2020-01-30 21:25:42 -08:00 committed by GitHub
parent 79f1756d8e
commit 1d79926d27
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 22 additions and 10 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
onnxruntime/core/providers/cpu/tensor/split.cc
View file

@ -18,6 +18,7 @@ ONNX_CPU_OPERATOR_VERSIONED_KERNEL(
std::vector<MLDataType>{
DataTypeImpl::GetTensorType<float>(),
DataTypeImpl::GetTensorType<int32_t>(),
DataTypeImpl::GetTensorType<int64_t>(),
DataTypeImpl::GetTensorType<std::string>()}),
Split);
@ -29,6 +30,7 @@ ONNX_CPU_OPERATOR_KERNEL(
std::vector<MLDataType>{
DataTypeImpl::GetTensorType<float>(),
DataTypeImpl::GetTensorType<int32_t>(),
DataTypeImpl::GetTensorType<int64_t>(),
DataTypeImpl::GetTensorType<std::string>()}),
Split);
@ -79,6 +81,8 @@ Status Split::Compute(OpKernelContext* context) const {
status = ComputeImpl<float>(*context, input);
else if (input.IsDataType<int32_t>())
status = ComputeImpl<int32_t>(*context, input);
else if (input.IsDataType<int64_t>())
status = ComputeImpl<int64_t>(*context, input);
else if (input.IsDataTypeString())
status = ComputeImpl<std::string>(*context, input);
else
@ -92,7 +96,7 @@ inline void copy_data(const T* src, T* dst, size_t count) {
memcpy(dst, src, count * sizeof(T));
}
template<>
template <>
inline void copy_data<std::string>(const std::string* src, std::string* dst, size_t count) {
const std::string* end = src + count;
std::copy(src, end, dst);

26
onnxruntime/test/providers/cpu/tensor/split_op_test.cc
View file

@ -11,7 +11,6 @@ template <class T>
using ShapeAndData = std::pair<const std::vector<int64_t>, const std::vector<T>>;
using ShapeAndFloatData = ShapeAndData<float>;
using ShapeAndInt32Data = ShapeAndData<int32_t>;
using ShapeAndStringData = ShapeAndData<std::string>;
using ExpectResult = OpTester::ExpectResult;
@ -65,16 +64,17 @@ TEST(SplitOperatorTest, Axis0EqualSplitFloat) {
RunTest<float>(axis, {}, input, outputs, false); //TensorRT parser: Assertion failed: axis != BATCH_DIM
}
TEST(SplitOperatorTest, Axis0EqualSplitInt32) {
template <typename T, typename = typename std::enable_if<std::is_integral<T>::value, T>::type>
static void SplitTestInt() {
const int64_t axis = 0;
std::vector<ShapeAndInt32Data> outputs;
std::vector<ShapeAndData<T>> outputs;
// input shape and data
ShapeAndInt32Data input = {{4, 2}, // shape
{1, 2,
3, 4,
5, 6,
7, 8}};
ShapeAndData<T> input = {{4, 2}, // shape
{1, 2,
3, 4,
5, 6,
7, 8}};
outputs.push_back({{2, 2},
{1, 2,
@ -84,7 +84,15 @@ TEST(SplitOperatorTest, Axis0EqualSplitInt32) {
{5, 6,
7, 8}});
RunTest<int32_t>(axis, {}, input, outputs, false); //TensorRT parser: Assertion failed: axis != BATCH_DIM
RunTest<T>(axis, {}, input, outputs, false); //TensorRT parser: Assertion failed: axis != BATCH_DIM
}
TEST(SplitOperatorTest, Axis0EqualSplitInt32) {
SplitTestInt<int32_t>();
}
TEST(SplitOperatorTest, Axis0EqualSplitInt64) {
SplitTestInt<int64_t>();
}
TEST(SplitOperatorTest, Axis0EqualSplitString) {