// 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;