def default_partition_fn(keys, shard_num):
"""The default partition function.
partition keys by "mod" strategy.
keys: a tensor presents the keys to be partitioned.
shard_num: the num of partitions
Returns:
a tensor with same shape as keys with type of `tf.int32`,
represents the corresponding partition-ids of keys.
"""
keys_op = ops.convert_to_tensor(keys, name="keys")
gpu_mode = pywrap_tensorflow.IsGoogleCudaEnabled()
with ops.colocate_with(keys_op):
if keys_op.dtype == dtypes.int64 and gpu_mode:
# This branch has low performance on some multi-CPU scenario,
# so we try to use default branch when GPUs are not available.
mask = constant_op.constant(0x7fffffff, dtypes.int64)
keys_int32 = math_ops.cast(bitwise_ops.bitwise_and(keys_op, mask),
dtypes.int32)
mod = math_ops.mod(keys_int32,
constant_op.constant(shard_num, dtypes.int32))
ids = math_ops.cast(mod, dtype=dtypes.int32)
elif keys_op.dtype == dtypes.string:
ids = string_ops.string_to_hash_bucket_fast(keys_op, shard_num)
mask = constant_op.constant(0x7fffffff, dtypes.int64)
ids = math_ops.cast(bitwise_ops.bitwise_and(ids, mask),
dtypes.int32)
else:
ids = math_ops.cast(math_ops.mod(keys_op, shard_num),
dtype=dtypes.int32)
return ids
def IsGoogleCudaEnabled():
return pywrap_tensorflow.IsGoogleCudaEnabled()
def _get_default_devices():
# use GPU if available
return ["/GPU:0"] if pywrap_tensorflow.IsGoogleCudaEnabled() else [
"/CPU:0",
]
