code: def testRngBitGenerator(self, algorithm, dtype):
      initial_state = array_ops.placeholder(np.uint64, shape=(2,))
      shape = (2, 3)
      res = xla.rng_bit_generator(
          algorithm, initial_state, shape, dtype=dtype
      )
  
      self.assertEqual(res[0].shape, initial_state.shape)
      self.assertEqual(res[1].shape, shape)
  
      # The initial_state has unknown dimension size
      initial_state = array_ops.placeholder(np.uint64, shape=(None,))
      shape = (2, 3)
      res = xla.rng_bit_generator(
          algorithm, initial_state, shape, dtype=dtype
      )
  
      self.assertEqual(
          res[0].shape.as_list(), initial_state.shape.as_list()
      )
      self.assertEqual(res[1].shape, shape)
  
      # The initial_state has unknown rank
      initial_state = array_ops.placeholder(np.uint64, shape=None)
      shape = (2, 3)
      res = xla.rng_bit_generator(
          algorithm, initial_state, shape, dtype=dtype
      )
  
      self.assertEqual(res[0].shape.as_list(), [None])
      self.assertEqual(res[1].shape, shape)
  
      # The output shape has unknown dimension
      initial_state = array_ops.placeholder(np.uint64, shape=(None,))
      shape = (None, 3)
      with self.assertRaisesRegex(
          TypeError, "Failed to convert elements .* to Tensor"
      ):
          res = xla.rng_bit_generator(
              algorithm, initial_state, shape, dtype=dtype
          )
  
imports: import functools
  from absl.testing import parameterized
  import numpy
  from tensorflow.compiler.tests import xla_test
  from tensorflow.compiler.tf2xla.ops import gen_xla_ops
  from tensorflow.compiler.tf2xla.python import xla
  from local_xla.xla import xla_data_pb2
  from tensorflow.python.eager import def_function
  from tensorflow.python.framework import dtypes
  from tensorflow.python.framework import errors
  from tensorflow.python.framework import function
  from tensorflow.python.framework import ops
  from tensorflow.python.framework import tensor_shape
  from tensorflow.python.framework import test_util
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import array_ops_stack
  from tensorflow.python.ops import random_ops_util
  from tensorflow.python.platform import googletest
label_desc: Use image node_id 16 for calling the XlaOpsShapeInferenceTest obj's underlying member method code with example usage: obj.testRngBitGenerator(algorithm, dtype) without return types