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

func.func @rng_bit_generator(%arg0: tensor<2xui64>) -> (tensor<2xui64>, tensor<10x12xui32>) {
  %output_state, %output = "vhlo.rng_bit_generator_v1"(%arg0) <{rng_algorithm = #vhlo<rng_algorithm_v1 DEFAULT>}> : (tensor<2xui64>) -> (tensor<2xui64>, tensor<10x12xui32>)
  func.return %output_state, %output : tensor<2xui64>, tensor<10x12xui32>
}

        rewriter.getContext(),
        *mlir::mhlo::symbolizeRngAlgorithm(xla_alg.value()));
    auto rng_bit_generator_op = rewriter.create<mhlo::RngBitGeneratorOp>(
        loc, op.getResultTypes(), algorithm_attr, op.getInitialState());

    rewriter.replaceOp(op, rng_bit_generator_op.getResults());

    return success();
  }
};

// Converts tf.XlaVariadicReduceV2 to mhlo.Reduce