| 52 | |
| 53 | template <typename T, typename Context> |
| 54 | void LinspaceKernel(const Context& dev_ctx, |
| 55 | const DenseTensor& start, |
| 56 | const DenseTensor& stop, |
| 57 | const DenseTensor& number, |
| 58 | DataType dtype, |
| 59 | DenseTensor* out) { |
| 60 | using XPUType = typename XPUTypeTrait<T>::Type; |
| 61 | T start_value = GetValueOfExpectedType<T, Context>(dev_ctx, start); |
| 62 | T stop_value = GetValueOfExpectedType<T, Context>(dev_ctx, stop); |
| 63 | int64_t num = GetValueOfExpectedType<int64_t, Context>(dev_ctx, number); |
| 64 | PADDLE_ENFORCE_GE(num, |
| 65 | 0, |
| 66 | common::errors::InvalidArgument( |
| 67 | "The num of linspace op should be larger " |
| 68 | "than or equal to 0, but received num is %d", |
| 69 | num)); |
| 70 | |
| 71 | out->Resize({num}); |
| 72 | T* out_data = dev_ctx.template Alloc<T>(out); |
| 73 | if (num == 0) { |
| 74 | return; |
| 75 | } |
| 76 | int r = xpu::linspace(dev_ctx.x_context(), |
| 77 | reinterpret_cast<XPUType*>(out_data), |
| 78 | static_cast<XPUType>(start_value), |
| 79 | static_cast<XPUType>(stop_value), |
| 80 | num); |
| 81 | PADDLE_ENFORCE_XDNN_SUCCESS(r, "linspace"); |
| 82 | } |
| 83 | |
| 84 | } // namespace phi |
| 85 | |