Check failure when running gen_stateless_random_ops.stateless_random_binomial

[dmc1778]

Click to expand!

### Issue Type Bug ### Have you reproduced the bug with TF nightly? Yes ### Source source ### Tensorflow Version 2.10.0 ### Custom Code Yes ### OS Platform and Distribution Ubuntu 22.04 ### Mobile device _No response_ ### Python version 3.9 ### Bazel version _No response_ ### GCC/Compiler version _No response_ ### CUDA/cuDNN version _No response_ ### GPU model and memory _No response_ ### Current Behaviour? ```shell I get check failure with the following input combination. ``` ### Standalone code to reproduce the issue ```shell import tensorflow as tf import os import numpy as np from tensorflow.python.ops import gen_stateless_random_ops try: shape_tensor = tf.random.uniform([1], minval=-256, maxval=257, dtype=tf.int64) shape = tf.identity(shape_tensor) seed_0 = 12 seed_1 = 34 seed = [seed_0,seed_1,] counts_tensor = tf.random.uniform([], dtype=tf.float32) counts = tf.identity(counts_tensor) probs_tensor = tf.random.uniform([], dtype=tf.float32) probs = tf.identity(probs_tensor) dtype = tf.float16 out = gen_stateless_random_ops.stateless_random_binomial(shape=shape,seed=seed,counts=counts,probs=probs,dtype=dtype,) except Exception as e: print("Error:"+str(e)) ``` ``` import tensorflow as tf import os import numpy as np from tensorflow.python.ops import stateless_random_ops try: shape_0 = 125091515651 shape = [shape_0,] seed_0 = 12 seed_1 = 34 seed = [seed_0,seed_1,] counts = 10.0 probs = 0.4 output_dtype = tf.float16 out = stateless_random_ops.stateless_random_binomial(shape=shape,seed=seed,counts=counts,probs=probs,output_dtype=output_dtype,) except Exception as e: print("Error:"+str(e)) ``` ### Relevant log output ```shell Error:{{function_node __wrapped__StatelessRandomBinomial_device_/job:localhost/replica:0/task:0/device:CPU:0}} Dimension -144 must be >= 0 [Op:StatelessRandomBinomial] 2023-01-07 13:26:54.833378: F tensorflow/core/framework/tensor.cc:719] Check failed: dtype() == expected_dtype (9 vs. 3) int32 expected, got int64 Aborted ``` ```

[tilakrayal]

@sachinprasadhs, I was able to reproduce the issue on tensorflow v2.11 and nightly. Kindly find the gist of it here.

[yongtang]

Added a PR #59233 for the fix.

[google-ml-butler[bot]]

Are you satisfied with the resolution of your issue? Yes No

[mihaimaruseac]

59300

[dmc1778]

59300

Also this one:

import tensorflow as tf
import numpy as np
from tensorflow.python.ops import gen_stateless_random_ops
try:
  try:
    with tf.device('/CPU'):
      shape_tensor = tf.random.uniform([1], minval=-256, maxval=257, dtype=tf.int64)
      shape = tf.identity(shape_tensor)
      seed_0 = 11
      seed_1 = 1
      seed = [seed_0,seed_1,]
      counts_tensor = tf.random.uniform([], dtype=tf.float32)
      counts = tf.identity(counts_tensor)
      probs_tensor = tf.random.uniform([], dtype=tf.float32)
      probs = tf.identity(probs_tensor)
      dtype = tf.float16
      out = gen_stateless_random_ops.stateless_random_binomial(shape=shape,seed=seed,counts=counts,probs=probs,dtype=dtype,)
  except Exception as e:
    print("Error:"+str(e))
  try:
    with tf.device('/GPU:0'):
      shape = tf.identity(shape_tensor)
      shape = tf.cast(shape, tf.int64)
      seed = [seed_0,seed_1,]
      counts = tf.identity(counts_tensor)
      counts = tf.cast(counts, tf.float32)
      probs = tf.identity(probs_tensor)
      probs = tf.cast(probs, tf.float32)
      dtype = tf.float16
      gen_stateless_random_ops.stateless_random_binomial(shape=shape,seed=seed,counts=counts,probs=probs,dtype=dtype,)
  except Exception as e:
    print("Error:"+str(e))
except Exception as e:
  print("Error:"+str(e))

[yongtang]

@mihaimaruseac do you know the details of the rollback (or the failed internal tests)? I will try to get a fix.

[mihaimaruseac]

@learning-to-play