REGISTER_KERNEL_BUILDER(Name("XlaEinsum").Device(DEVICE),                    \
                          XlaCompileOnDemandOp);                               \
  REGISTER_KERNEL_BUILDER(Name("XlaRngBitGenerator")                           \
                              .HostMemory("algorithm")                         \
                              .HostMemory("shape")                             \

                                         "XlaReduce",
                                         "XlaReduceWindow",
                                         "XlaReplicaId",
                                         "XlaRngBitGenerator",
                                         "XlaScatter",
                                         "XlaSelectAndScatter",
                                         "XlaSelfAdjointEig",

      "XlaRecv",
      "XlaReduce",
      "XlaReducePrecision",
      "XlaReduceScatter",
      "XlaReduceWindow",
      "XlaRemoveDynamicDimensionSize",
      "XlaReplicaId",
      "XlaRngBitGenerator",
      "XlaScatter",
      "XlaSelectAndScatter",
      "XlaSelfAdjointEig",
      "XlaSend",
      "XlaSetBound",
      "XlaSetDynamicDimensionSize",
      "XlaSharding",
      "XlaSort",

  %output = "tf.XlaSort"(%input): (tensor<2x3xi64>) -> (tensor<2x3xi64>)
  func.return %output : tensor<2x3xi64>
}

//===----------------------------------------------------------------------===//
// tf.XlaRngBitGenerator legalization
//===----------------------------------------------------------------------===//

// CHECK-LABEL: @xla_rng_bit_generator
// CHECK-SAME: %[[STATE:.*]]: tensor<2xui64>

    return xla::RandomAlgorithm::RNG_THREE_FRY;
  } 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;

    TF_Int32Tensor:$id
  );

  // Constant folding is disabled for this op as it is a runtime op and can't
  // constant folded at the compile time.
}

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

  let description = [{
Wraps the XLA RngBitGenerator operator, documented at