Fix Pad and Gather incorrect usage of HasType helpers. (#7146)

This commit is contained in:
Edward Chen 2021-03-26 17:36:31 -07:00 committed by GitHub
parent ab86634c36
commit 63d9d5afd3
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 7 additions and 14 deletions

View file

@ -148,12 +148,12 @@ Status Gather::Compute(OpKernelContext* context) const {
concurrency::ThreadPool* tp = context->GetOperatorThreadPool();
if (utils::HasTypeWithSameSize<EnabledIndexTypes, int32_t>() &&
if (utils::HasType<EnabledIndexTypes, int32_t>() &&
p.indices_tensor->IsDataType<int32_t>()) {
return GatherCopyData<int32_t>(p.indices_tensor, src_base, dst_base, is_string_type, element_bytes,
block_size, M, N, data_batch_bytes, gathered_batch_bytes, input_data_shape, p.axis, tp);
}
if (utils::HasTypeWithSameSize<EnabledIndexTypes, int64_t>() &&
if (utils::HasType<EnabledIndexTypes, int64_t>() &&
p.indices_tensor->IsDataType<int64_t>()) {
return GatherCopyData<int64_t>(p.indices_tensor, src_base, dst_base, is_string_type, element_bytes,
block_size, M, N, data_batch_bytes, gathered_batch_bytes, input_data_shape, p.axis, tp);

View file

@ -244,7 +244,7 @@ static Status PadImpl(OpKernelContext* ctx,
const std::vector<int64_t>& slices,
const Mode& mode,
T value) {
if (!utils::HasType<AllEnabledPadTypes, T>()) {
if (!utils::HasTypeWithSameSize<AllEnabledPadTypes, T>()) {
return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Input data type not supported in this build.");
}
@ -517,22 +517,15 @@ Status Pad::Compute(OpKernelContext* ctx) const {
slices_to_use = &slices_;
}
Status pad_status{};
switch (element_size) {
case sizeof(uint32_t):
pad_status = PadImpl<uint32_t>(ctx, *pads_to_use, *slices_to_use, mode_, value.u32);
break;
return PadImpl<uint32_t>(ctx, *pads_to_use, *slices_to_use, mode_, value.u32);
case sizeof(uint64_t):
pad_status = PadImpl<uint64_t>(ctx, *pads_to_use, *slices_to_use, mode_, value.u64);
break;
return PadImpl<uint64_t>(ctx, *pads_to_use, *slices_to_use, mode_, value.u64);
case sizeof(uint8_t):
pad_status = PadImpl<uint8_t>(ctx, *pads_to_use, *slices_to_use, mode_, value.u8);
break;
return PadImpl<uint8_t>(ctx, *pads_to_use, *slices_to_use, mode_, value.u8);
default:
pad_status = ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Unsupported input data type of ", data_type);
break;
return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Unsupported input data type of ", data_type);
}
return pad_status;
}
}; // namespace onnxruntime