A check fail can be triggered in LearnedUnigramCandidateSampler

[shijy16]

Click to expand!

Issue Type

Bug

Have you reproduced the bug with TF nightly?

Yes

Source

binary

Tensorflow Version

2.13.0-dev20230331

Custom Code

No

OS Platform and Distribution

Linux Ubuntu 20.04

Mobile device

No response

Python version

3.10

Bazel version

No response

GCC/Compiler version

No response

CUDA/cuDNN version

CUDA 11.8

GPU model and memory

No response

Current Behaviour?

The following code can trigger a crash in `tf.raw_ops.LearnedUnigramCandidateSampler` due to check-fail in the latest version of TensorFlow.

Standalone code to reproduce the issue

import tensorflow as tf
with tf.device("GPU:0"):
    num_true = 13
    num_sampled = 48
    unique = True
    range_max = 3031324185113192368
    seed = 93
    seed2 = 11
    true_classes = tf.saturate_cast(tf.random.uniform([14, 13], minval=-1024, maxval=1024, dtype=tf.int64), dtype=tf.int64)
    res = tf.raw_ops.LearnedUnigramCandidateSampler(
        num_true=num_true,
        num_sampled=num_sampled,
        unique=unique,
        range_max=range_max,
        seed=seed,
        seed2=seed2,
        true_classes=true_classes,
    )

Relevant log output

2023-04-01 16:20:32.160750: I tensorflow/core/util/port.cc:110] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.
2023-04-01 16:20:32.211959: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
To enable the following instructions: AVX2 AVX512F AVX512_VNNI FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
2023-04-01 16:20:33.026789: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT
2023-04-01 16:20:34.550122: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1635] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 14561 MB memory:  -> device: 0, name: Tesla V100-PCIE-16GB, pci bus id: 0000:2f:00.0, compute capability: 7.0
2023-04-01 16:20:39.689581: F tensorflow/core/kernels/range_sampler.cc:183] Check failed: range < kint32max (3031324185113192368 vs. 2147483647)
Aborted (core dumped)

[tilakrayal]

@shijy16,
I can see the range_max = 3031324185113192368 which you are trying is a very big value. Could you please provide any usecase you are trying the big value.

range_max	An int that is >= 1. The sampler will sample integers from the interval [0, range_max).

Thank you!

[shijy16]

@tilakrayal
Hi, I am doing API testing to find potential bugs/vulnerabilites.
Although the error message seems reasonable, the abort/crash is unacceptable. Instead, an exception that can be catched with external code like try: ... except: ... should be placed here.

[mihaimaruseac]

@shijy16 Please don't file vulnerabilities on GitHub. Please consult https://github.com/tensorflow/tensorflow/blob/master/SECURITY.md and follow rules for responsible disclosure.