// CHECK-NEXT:  return %0 : tensor<2x3x2xi32>
// CHECK-NEXT:}

func.func @rng_bit_generator(%arg0: tensor<2xui64>) -> (tensor<2xui64>, tensor<10x12xui32>) {

  // CHECK-NEXT: %1 = mhlo.constant dense<3> : tensor<i32>
  %cst_0 = "tf.Const"() {value = dense<3> : tensor<i32>} : () -> tensor<i32>
  // CHECK-NEXT: %[[OUTPUT_STATE:.*]], %[[OUTPUT:.*]] = "mhlo.rng_bit_generator"(%[[STATE]]) <{rng_algorithm = #mhlo.rng_algorithm<DEFAULT>}> : (tensor<2xui64>) -> (tensor<2xui64>, tensor<10x12xui32>)
  // CHECK-NEXT: return %[[OUTPUT_STATE]], %[[OUTPUT]] : tensor<2xui64>, tensor<10x12xui32>

  } else if (alg == tensorflow::RNG_ALG_AUTO_SELECT) {
    return xla::RandomAlgorithm::RNG_DEFAULT;
  }
  return std::nullopt;
}

// Converts tf.XlaRngBitGenerator op to mhlo.RngBitGenerator op.
class ConvertXlaRngBitGeneratorOp
    : public OpRewritePattern<TF::XlaRngBitGeneratorOp> {
 public:
  using OpRewritePattern::OpRewritePattern;

def TF_XlaRngBitGeneratorOp : TF_Op<"XlaRngBitGenerator", [Pure]> {
  let summary = "Stateless PRNG bit generator.";

  let description = [{
Wraps the XLA RngBitGenerator operator, documented at
 https://www.tensorflow.org/performance/xla/operation_semantics#rngbitgenerator.
  }];

  let arguments = (ins
    Arg<TF_Int32Tensor, [{The PRNG algorithm to use, one of