// CHECK-SAME: ->{{.*}}mhlo.sharding = "\04\05\06"{{.*}}mhlo.sharding = "\01\02\03"
func.func @func(%arg0: tensor<*xi32> {tf.aliasing_output = 1 : i64},
           %arg1: tensor<*xi32> {tf.aliasing_output = 0 : i64}) -> (tensor<*xi32>, tensor<*xi32>) {
  %0 = "tf.XlaSharding"(%arg0) { _XlaSharding = "\01\02\03"} : (tensor<*xi32>) -> tensor<*xi32>

    %3 = "tf.XlaCallModule"(%2, %1, %0) <{Sout = [#tf_type.shape<?x2>], module = "", version = 9 : i64}> {_entry_function = @composite_dot_general_with_bias_and...

      %3 = stablehlo.maximum %arg1, %arg2 : tensor<f32>
      stablehlo.return %3 : tensor<f32>
  }) {
    window_dimensions = array<i64: 1, 1, 2, 2>,
    window_strides = array<i64: 1, 1, 2, 2>,
    base_dilations = array<i64: 1, 1, 2, 2>,
    window_dilations = array<i64: 1, 1, 2, 2>
  } : (tensor<1x4x16x16xf32>, tensor<f32>) -> tensor<1x4x15x15xf32>
  return %2 : tensor<1x4x15x15xf32>
}

  %5 = stablehlo.convolution(%arg0, %0) dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], window = {pad = [[1, 1], [1, 1]]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64} : (tensor<?x28x28x1xf32>, tensor<3x3x1x16xf32>) -> tensor<?x28x28x16xf32>

  // CHECK: %[[dq2:.*]] = "quantfork.dcast"(%[[q2]])
  // CHECK: stablehlo.convert %[[dq2]]
  %0 = "quantfork.stats"(%arg0) {bitsNum = 8 : i64, layerStats = dense<[-0.999415695, 0.99998933]> : tensor<2xf32>, narrowRange = false} : (tensor<2x3xf32>) -> tensor<2x3xf32>
  %1 = "quantfork.stats"(%0) {bitsNum = 8 : i64, layerStats = dense<[-2.0, 2.0]> : tensor<2xf32>, narrowRange = false} : (tensor<2x3xf32>) -> tensor<2x3xf32>

  %batch = "tfr.constant_tensor"(%batch_size) : (i64) -> tensor<i64>
  %batch1 = "tfr.cast"(%batch) : (tensor<i64>) -> !tfr.tensor
  %calls = "tfr.constant_tensor"(%num_parallel_calls) : (i64) -> tensor<i64>
  %calls1 = "tfr.cast"(%calls) : (tensor<i64>) -> !tfr.tensor
  %drop = "tfr.constant_tensor"(%drop_remainder) : (i1) -> tensor<i1>

}

// CHECK-LABEL: @tf__my_max_pool
tfr.func @tf__my_max_pool(%input_: !tfr.tensor, %stride_w: i64{tfr.name="stride_w"}, %stride_h: i64{tfr.name="stride_h"}) -> (!tfr.tensor) {
  %cst_1 = arith.constant 1 : i64
  %stride = "tfr.build_list"(%cst_1, %stride_w, %stride_h, %cst_1) : (i64, i64, i64, i64) -> !tfr.attr
  %filter = tfr.constant [1, 2, 2, 1] -> !tfr.attr
  %padding = tfr.constant "VALID" -> !tfr.attr

module {
func.func @main(%arg0: tensor<3x3xf32>, %arg1: tensor<3x3xf32>) -> tensor<6x3xf32> {
  %axis = "tf.Const"() { value = dense<0> : tensor<i64> } : () -> tensor<i64>
  %1 = "tf.ConcatV2"(%arg0, %arg1, %axis) : (tensor<3x3xf32>, tensor<3x3xf32>, tensor<i64>) -> tensor<6x3xf32>
  func.return %1 : tensor<6x3xf32>
}

	i16 := types.Types[types.TINT16]
	i32 := types.Types[types.TINT32]
	i64 := types.Types[types.TINT64]
	f32 := types.Types[types.TFLOAT32]
	f64 := types.Types[types.TFLOAT64]
	c64 := types.Types[types.TCOMPLEX64]
	c128 := types.Types[types.TCOMPLEX128]
	ft := mkFuncType(nil,
		[]*types.Type{
			i8, i16, i32, i64,
			f32, f32, f64, f64,
			i8, i16, i32, i64,
			f32, f32, f64, f64,
			c128, c128, c128, c128, c64,

    appended at the end).
  }];

  let arguments = (ins
    TFL_TensorOf<[I32, I64]>:$indices,
    TFL_I32Tensor:$depth,
    TFL_TensorOf<[F32, I32, I64, I1, I8, UI8]>:$on_value,
    TFL_TensorOf<[F32, I32, I64, I1, I8, UI8]>:$off_value,

    I32Attr:$axis
  );

  let results = (outs
    TFL_TensorOf<[F32, I32, I64, I1, I8, UI8]>:$output
  );

  let hasOptions = 1;
}