// CHECK: return %[[UNIFORM_QUANTIZE_1]] : tensor<1x3x!quant.uniform<i8:f32, {{.*}}>>

// CHECK-PER-TENSOR: func.func private @quantized_dot_general_fn(%[[ARG_1:.+]]: tensor<1x2x!quant.uniform<i8:f32, {{.*}}>>, %[[ARG_2:.+]]: tensor<2x3x!quant.uniform<i8<-127:127>:f32, {{.*}}>>) -> tensor<1x3x!quant.uniform<i8:f32, {{.*}}>> attributes {_from_xla_call_module}

  %1 = "tfl.dequantize"(%0) : (tensor<2x1x!quant.uniform<i16:f32, 1.0>>) -> (tensor<2x1xf32>)
  %2 = "tfl.quantize"(%arg1) {qtype = tensor<2x3x!quant.uniform<i16:f32, 1.0>>} : (tensor<2x3xf32>) -> tensor<2x3x!quant.uniform<i16:f32, 1.0>>
  %3 = "tfl.dequantize"(%2) : (tensor<2x3x!quant.uniform<i16:f32, 1.0>>) -> (tensor<2x3xf32>)

// CHECK-DAG: %[[input_14:.*]] = "tfl.dequantize"({{.*}}) : (tensor<2x!quant.uniform<i32:f32, 1.1827185827747504E-4>>) -> tensor<2xf32>
// CHECK-DAG: %[[input_15:.*]] = "tfl.dequantize"({{.*}}) : (tensor<2x!quant.uniform<i32:f32, 1.8229184289320815E-4>>) -> tensor<2xf32>
// CHECK-DAG: %[[input_16:.*]] = "tfl.dequantize"({{.*}}) : (tensor<4x2x!quant.uniform<i8<-127:127>:f32, 0.013545581674951268>>) -> tensor<4x2xf32>

  func.return %5: tensor<2x1x!quant.uniform<i8:f32, 0.003:-128>>
}

// CHECK:   func @quantizedOpOnly(%[[VAL_0:.*]]: tensor<1x!quant.uniform<i8:f32, 3.000000e-03:-128>>, %[[VAL_1:.*]]: tensor<1x!quant.uniform<i8:f32, 3.000000e-03:-128>>) -> tensor<2x1x!quant.uniform<i8:f32, 3.000000e-03:-128>> {

// ```
//
// The rewritten pattern looks like:
//
// ```
// %1 = stablehlo.uniform_quantize %0  // Input f32 -> uniform quantized type.
// %2 = stablehlo.constant  // Filter uniform quantized type.
// %3 = stablehlo.convolution(%1, %2)   // In uniform quantized type.
// %4 = stablehlo.uniform_dequantize %3  // Dequantize the output.
// ```
class ComposeUniformQuantizedConvolutionOp

    rng = np.random.default_rng(seed=42)
    input_data = ops.convert_to_tensor(
        rng.uniform(low=0.0, high=1.0, size=static_input_shape).astype(
            np.float32
        )
    )

    def data_gen() -> repr_dataset.RepresentativeDataset:
      for _ in range(100):
        yield {
            'input_tensor': rng.uniform(
                low=0.0, high=1.0, size=static_input_shape
            ).astype(np.float32)

^bb0(%arg0: tensor<1x2xf32>):
  %cst_0 = arith.constant dense<[1, 0]> : tensor<2xi32>
  %0 = "tfl.quantize"(%arg0){qtype = tensor<1x2x!quant.uniform<u8:f32, 1.0>>}: (tensor<1x2xf32>) -> (tensor<1x2x!quant.uniform<u8:f32, 1.0>>)
  %1 = "tfl.dequantize"(%0): (tensor<1x2x!quant.uniform<u8:f32, 1.0>>) -> (tensor<1x2xf32>)
  %2 = "tf.Transpose"(%1, %cst_0): (tensor<1x2xf32>, tensor<2xi32>) -> tensor<2x1xf32>
  func.return %2 : tensor<2x1xf32>

        self.embedding_w = np.minimum(np.maximum(self.embedding_w, -4), 4)

        self.conv_filters = np.random.uniform(
            low=-10, high=10, size=filter_shape
        ).astype('f4')

        second_conv_filter_shape = (3, 3, filter_shape[-1], 1)
        self.second_conv_filters = np.random.uniform(
            low=-10, high=10, size=second_conv_filter_shape
        ).astype('f4')

      @def_function.function(

  }];
}

def TF_StatefulUniformFullIntOp : TF_Op<"StatefulUniformFullInt", []> {
  let summary = "Outputs random integers from a uniform distribution.";

  let description = [{
The generated values are uniform integers covering the whole range of `dtype`.
  }];

  let arguments = (ins
    Arg<TF_ResourceTensor, "", [TF_VariableRead,TF_VariableWrite]>:$resource,

	uc.finalSeats = uint64(finalSeats)
	uc.padDuration = duration
	return uc
}

// uniformScenario describes a scenario based on the given set of uniform clients.
// Each uniform client specifies a number of threads, each of which alternates between thinking
// and making a synchronous request through the QueueSet.
// The test measures how much concurrency each client got, on average, over