|
|
|
@ -1,5 +1,5 @@
|
|
|
|
|
/**
|
|
|
|
|
* Copyright 2020 Huawei Technologies Co., Ltd
|
|
|
|
|
* Copyright 2020-2021 Huawei Technologies Co., Ltd
|
|
|
|
|
*
|
|
|
|
|
* Licensed under the Apache License, Version 2.0 (the "License");
|
|
|
|
|
* you may not use this file except in compliance with the License.
|
|
|
|
@ -18,6 +18,16 @@
|
|
|
|
|
|
|
|
|
|
namespace mindspore {
|
|
|
|
|
namespace kernel {
|
|
|
|
|
MS_REG_GPU_KERNEL_TWO(
|
|
|
|
|
UnsortedSegmentSum,
|
|
|
|
|
KernelAttr().AddInputAttr(kNumberTypeFloat64).AddInputAttr(kNumberTypeInt32).AddOutputAttr(kNumberTypeFloat64),
|
|
|
|
|
UnsortedSegmentSumGpuKernel, double, int)
|
|
|
|
|
|
|
|
|
|
MS_REG_GPU_KERNEL_TWO(
|
|
|
|
|
UnsortedSegmentSum,
|
|
|
|
|
KernelAttr().AddInputAttr(kNumberTypeFloat64).AddInputAttr(kNumberTypeInt64).AddOutputAttr(kNumberTypeFloat64),
|
|
|
|
|
UnsortedSegmentSumGpuKernel, double, int64_t)
|
|
|
|
|
|
|
|
|
|
MS_REG_GPU_KERNEL_TWO(
|
|
|
|
|
UnsortedSegmentSum,
|
|
|
|
|
KernelAttr().AddInputAttr(kNumberTypeFloat32).AddInputAttr(kNumberTypeInt32).AddOutputAttr(kNumberTypeFloat32),
|
|
|
|
@ -39,6 +49,36 @@ MS_REG_GPU_KERNEL_TWO(
|
|
|
|
|
UnsortedSegmentSumGpuKernel, int, int64_t)
|
|
|
|
|
|
|
|
|
|
// Re-registration with 3 inputs - dynamic shape mode - sets of Int64/Int32 num segments types
|
|
|
|
|
MS_REG_GPU_KERNEL_TWO(UnsortedSegmentSum,
|
|
|
|
|
KernelAttr()
|
|
|
|
|
.AddInputAttr(kNumberTypeFloat64)
|
|
|
|
|
.AddInputAttr(kNumberTypeInt32)
|
|
|
|
|
.AddInputAttr(kNumberTypeInt64)
|
|
|
|
|
.AddOutputAttr(kNumberTypeFloat64),
|
|
|
|
|
UnsortedSegmentSumGpuKernel, double, int)
|
|
|
|
|
MS_REG_GPU_KERNEL_TWO(UnsortedSegmentSum,
|
|
|
|
|
KernelAttr()
|
|
|
|
|
.AddInputAttr(kNumberTypeFloat64)
|
|
|
|
|
.AddInputAttr(kNumberTypeInt32)
|
|
|
|
|
.AddInputAttr(kNumberTypeInt32)
|
|
|
|
|
.AddOutputAttr(kNumberTypeFloat64),
|
|
|
|
|
UnsortedSegmentSumGpuKernel, double, int)
|
|
|
|
|
|
|
|
|
|
MS_REG_GPU_KERNEL_TWO(UnsortedSegmentSum,
|
|
|
|
|
KernelAttr()
|
|
|
|
|
.AddInputAttr(kNumberTypeFloat64)
|
|
|
|
|
.AddInputAttr(kNumberTypeInt64)
|
|
|
|
|
.AddInputAttr(kNumberTypeInt64)
|
|
|
|
|
.AddOutputAttr(kNumberTypeFloat64),
|
|
|
|
|
UnsortedSegmentSumGpuKernel, double, int64_t)
|
|
|
|
|
MS_REG_GPU_KERNEL_TWO(UnsortedSegmentSum,
|
|
|
|
|
KernelAttr()
|
|
|
|
|
.AddInputAttr(kNumberTypeFloat64)
|
|
|
|
|
.AddInputAttr(kNumberTypeInt64)
|
|
|
|
|
.AddInputAttr(kNumberTypeInt32)
|
|
|
|
|
.AddOutputAttr(kNumberTypeFloat64),
|
|
|
|
|
UnsortedSegmentSumGpuKernel, double, int64_t)
|
|
|
|
|
|
|
|
|
|
MS_REG_GPU_KERNEL_TWO(UnsortedSegmentSum,
|
|
|
|
|
KernelAttr()
|
|
|
|
|
.AddInputAttr(kNumberTypeFloat32)
|
|
|
|
|