experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Add GPU support for complex types in OneHot

Browse Source 
PiperOrigin-RevId: 314621230
Change-Id: I23bb95ff53cc2e8b16d73cbc1c22f42af6461170
This commit is contained in:
Gaurav Jain 2020-06-03 15:55:39 -07:00 committed by TensorFlower Gardener
parent 35875e5b96
commit cfc9a6852e
2 changed files with 6 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
tensorflow/core/kernels/one_hot_op.cc
View File

@ -164,6 +164,8 @@ TF_CALL_GPU_NUMBER_TYPES(DECLARE_GPU_SPEC);
TF_CALL_bool(DECLARE_GPU_SPEC);
TF_CALL_int32(DECLARE_GPU_SPEC);
TF_CALL_int64(DECLARE_GPU_SPEC);
TF_CALL_complex64(DECLARE_GPU_SPEC);
TF_CALL_complex128(DECLARE_GPU_SPEC);
#undef DECLARE_GPU_SPEC_INDEX
#undef DECLARE_GPU_SPEC
@ -188,6 +190,8 @@ TF_CALL_GPU_NUMBER_TYPES(REGISTER_ONE_HOT_GPU);
TF_CALL_bool(REGISTER_ONE_HOT_GPU);
TF_CALL_int32(REGISTER_ONE_HOT_GPU);
TF_CALL_int64(REGISTER_ONE_HOT_GPU);
TF_CALL_complex64(REGISTER_ONE_HOT_GPU);
TF_CALL_complex128(REGISTER_ONE_HOT_GPU);
#undef REGISTER_ONE_HOT_GPU_INDEX
#undef REGISTER_ONE_HOT_GPU

2
tensorflow/core/kernels/one_hot_op_gpu.cu.cc
View File

@ -41,6 +41,8 @@ TF_CALL_GPU_NUMBER_TYPES(DEFINE_GPU_SPEC);
TF_CALL_bool(DEFINE_GPU_SPEC);
TF_CALL_int32(DEFINE_GPU_SPEC);
TF_CALL_int64(DEFINE_GPU_SPEC);
TF_CALL_complex64(DEFINE_GPU_SPEC);
TF_CALL_complex128(DEFINE_GPU_SPEC);
#undef DEFINE_GPU_SPEC_INDEX
#undef DEFINE_GPU_SPEC