diff options
| author | Francisco Massa <fvsmassa@gmail.com> | 2017-07-19 21:37:32 +0300 |
|---|---|---|
| committer | Soumith Chintala <soumith@gmail.com> | 2017-07-19 22:43:00 +0300 |
| commit | ea2cced1a95de3d6d3ac3b844e9ad2910efebce2 (patch) | |
| tree | 66bc07149628804dffa83c418794a585de235613 /lib | |
| parent | 07939f76d5259db7d6a014430bbb599109343fd7 (diff) | |
Add CUDA support for arange
Also enables CUDA for range
Diffstat (limited to 'lib')
| -rw-r--r-- | lib/THC/generic/THCTensorMath.cu | 11 | ||||
| -rw-r--r-- | lib/THC/generic/THCTensorMath.h | 1 |
2 files changed, 12 insertions, 0 deletions
diff --git a/lib/THC/generic/THCTensorMath.cu b/lib/THC/generic/THCTensorMath.cu index 0eed5a9..d47102c 100644 --- a/lib/THC/generic/THCTensorMath.cu +++ b/lib/THC/generic/THCTensorMath.cu @@ -449,4 +449,15 @@ void THCTensor_(range)(THCState *state, THCTensor *r_, accreal xmin, accreal xma THCudaCheck(cudaGetLastError()); } +void THCTensor_(arange)(THCState* state, THCTensor *r_, accreal xmin, accreal xmax, accreal step) { +#if defined(THC_REAL_IS_FLOAT) || defined(THC_REAL_IS_DOUBLE) || defined(THC_REAL_IS_HALF) + int m = fmod(xmax - xmin, step) == 0; +#else + int m = (xmax - xmin) % step == 0; +#endif + if (m) + xmax -= step; + THCTensor_(range)(state, r_, xmin, xmax, step); +} + #endif diff --git a/lib/THC/generic/THCTensorMath.h b/lib/THC/generic/THCTensorMath.h index aae6775..c8fb35b 100644 --- a/lib/THC/generic/THCTensorMath.h +++ b/lib/THC/generic/THCTensorMath.h @@ -26,5 +26,6 @@ THC_API void THCTensor_(logspace)(THCState *state, THCTensor *r_, real a, real b #endif THC_API void THCTensor_(range)(THCState *state, THCTensor *r_, accreal xmin, accreal xmax, accreal step); +THC_API void THCTensor_(arange)(THCState *state, THCTensor *r_, accreal xmin, accreal xmax, accreal step); #endif |