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

Add uint64 support to DynamicPartition and FloorMod on CPU.

Browse Source 
PiperOrigin-RevId: 301389554
Change-Id: I5414bc9f35e7f42a52c63767aa23f2508dcb3e37
This commit is contained in:
A. Unique TensorFlower 2020-03-17 09:30:23 -07:00 committed by TensorFlower Gardener
parent 2e62c905e4
commit 86dcc9732d
3 changed files with 5 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
tensorflow/core/kernels/cwise_op_floor_mod.cc
View File

@ -16,7 +16,8 @@ limitations under the License.
#include "tensorflow/core/kernels/cwise_ops_common.h" #include "tensorflow/core/kernels/cwise_ops_common.h"
namespace tensorflow { namespace tensorflow {
REGISTER2(BinaryOp, CPU, "FloorMod", functor::safe_floor_mod, int32, int64); REGISTER3(BinaryOp, CPU, "FloorMod", functor::safe_floor_mod, int32, int64,
uint64);
REGISTER2(BinaryOp, CPU, "FloorMod", functor::floor_fmod, float, double); REGISTER2(BinaryOp, CPU, "FloorMod", functor::floor_fmod, float, double);
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM

2
tensorflow/core/kernels/dynamic_partition_op.cc
View File

@ -164,6 +164,8 @@ class DynamicPartitionOp : public DynamicPartitionOp_Shared {
DynamicPartitionOp<T>) DynamicPartitionOp<T>)
TF_CALL_ALL_TYPES(REGISTER_DYNAMIC_PARTITION); TF_CALL_ALL_TYPES(REGISTER_DYNAMIC_PARTITION);
// For partitioning fingerprints.
TF_CALL_uint64(REGISTER_DYNAMIC_PARTITION);
#undef REGISTER_DYNAMIC_PARTITION #undef REGISTER_DYNAMIC_PARTITION
} // namespace tensorflow } // namespace tensorflow

2
tensorflow/core/ops/math_ops.cc
View File

@ -587,7 +587,7 @@ REGISTER_OP("FloorMod")
.Input("x: T") .Input("x: T")
.Input("y: T") .Input("y: T")
.Output("z: T") .Output("z: T")
.Attr("T: {int32, int64, bfloat16, half, float, double}") .Attr("T: {int32, int64, uint64, bfloat16, half, float, double}")
.SetShapeFn(shape_inference::BroadcastBinaryOpShapeFn); .SetShapeFn(shape_inference::BroadcastBinaryOpShapeFn);
REGISTER_OP("TruncateMod") REGISTER_OP("TruncateMod")