type PCG struct {
	hi uint64
	lo uint64
}

// NewPCG returns a new PCG seeded with the given values.
func NewPCG(seed1, seed2 uint64) *PCG {
	return &PCG{seed1, seed2}
}

// Seed resets the PCG to behave the same way as NewPCG(seed1, seed2).
func (p *PCG) Seed(seed1, seed2 uint64) {
	p.hi = seed1
	p.lo = seed2
}

// MarshalBinary implements the encoding.BinaryMarshaler interface.
func (p *PCG) MarshalBinary() ([]byte, error) {

)

func wyhash(key []byte, seed, len uint64) uint64 {
	p := key
	i := len
	var a, b uint64
	seed ^= m1

	if i > 16 {
		if i > 48 {
			seed1 := seed
			seed2 := seed
			for ; i > 48; i -= 48 {
				seed = mix(r8(p)^m2, r8(p[8:])^seed)
				seed1 = mix(r8(p[16:])^m3, r8(p[24:])^seed1)
				seed2 = mix(r8(p[32:])^m4, r8(p[40:])^seed2)
				p = p[48:]
			}
			seed ^= seed1 ^ seed2
		}

  %4 = "tf.Multinomial"(%0, %1) {device = "/job:localhost/replica:0/task:0/device:CPU:0", seed = 0 : i64, seed2 = 0 : i64} : (tensor<*xf32>, tensor<*xi32>) -> tensor<*xi64>
  %5 = "tf.Multinomial"(%0, %1) {device = "/job:localhost/replica:0/task:0/device:CPU:0", seed = 0 : i64, seed2 = 0 : i64} : (tensor<*xf32>, tensor<*xi32>) -> tensor<*xi64>

	default:
		l := s
		if l > 48 {
			seed1 := seed
			seed2 := seed
			for ; l > 48; l -= 48 {
				seed = mix(r8(p)^hashkey[1], r8(add(p, 8))^seed)
				seed1 = mix(r8(add(p, 16))^hashkey[2], r8(add(p, 24))^seed1)
				seed2 = mix(r8(add(p, 32))^hashkey[3], r8(add(p, 40))^seed2)
				p = add(p, 48)
			}
			seed ^= seed1 ^ seed2
		}
		for ; l > 16; l -= 16 {
			seed = mix(r8(p)^hashkey[1], r8(add(p, 8))^seed)

	for n := b.N; n > 0; n-- {
		t += p.Uint64()
	}
	Sink = t
}

func TestPCGMarshal(t *testing.T) {
	var p PCG
	const (
		seed1 = 0x123456789abcdef0
		seed2 = 0xfedcba9876543210
		want  = "pcg:\x12\x34\x56\x78\x9a\xbc\xde\xf0\xfe\xdc\xba\x98\x76\x54\x32\x10"
	)
	p.Seed(seed1, seed2)
	data, err := p.MarshalBinary()
	if string(data) != want || err != nil {
		t.Errorf("MarshalBinary() = %q, %v, want %q, nil", data, err, want)

def LegalizeRandomUniform : Pat<
  (TF_RandomUniformOp $shape, $seed, $seed2),
  (TFL_RandomUniformOp $shape, (convertIntAttrTo64Bit $seed),
   (convertIntAttrTo64Bit $seed2))>;

def LegalizeRandomStandardNormal : Pat<
  (TF_RandomStandardNormalOp $shape, $seed, $seed2),
  (TFL_RandomStandardNormalOp $shape, (convertIntAttrTo64Bit $seed),
   (convertIntAttrTo64Bit $seed2))>;

def LegalizeMultinomial : Pat<

func.func @random_uniform_simple(%arg0: tensor<3xi32>) -> tensor<12x?x64xf32> {
  // expected-remark@+1 {{lowering requires operand #0 to be a constant}}
  %0 = "tf.RandomUniform"(%arg0) {device = "", seed = 0 : i64, seed2 = 0 : i64} : (tensor<3xi32>) -> tensor<12x?x64xf32>
  func.return %0 : tensor<12x?x64xf32>
}

// -----
module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {

    // CHECK-NOT: tf.RandomUniformInt
    %1 = "tf.RandomUniformInt"(%0, %arg0, %arg1) {seed = 0 : i64, seed2 = 0 : i64} : (tensor<1xi32>, tensor<i32>, tensor<i32>) -> tensor<1000xi32>
    func.return %1 : tensor<1000xi32>
  }

  // CHECK-LABEL: multinomial
  func.func @multinomial(%arg0: tensor<2x4xf32>, %seed: tensor<i32>, %seed2: tensor<i32>) -> tensor<2x10xi32> {
    // CHECK-NOT: tf.Multinomial

  // Force consistent pooling regions for unit testing.
  auto y = FractionalAvgPool(
      scope_, x, {1, 1.2, 1.9, 1},
      FractionalAvgPool::Deterministic(true).Overlapping(true).Seed(1).Seed2(
          2));
  TensorShape y_shape({1, 2, 3, 1});
  RunTest(x, x_shape, y.output, y_shape);
}

TEST_F(NNGradTest, FractionalMaxPoolGradHelper) {
  TensorShape x_shape({1, 3, 7, 1});

//===----------------------------------------------------------------------===//
// TODO(b/148269299): handle random number generator seeds/states correctly.

class MHLO_RngDistributionValue<string enumStr> :
  ConstantAttr<MHLO_RngDistributionAttr, "::mlir::mhlo::RngDistribution::" # enumStr>;

def : Pat<(TF_RandomUniformOp:$old $shape, $seed, $seed2),
          (MHLO_RngOp
            (MHLO_ConstantOp