// 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>
}

  %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>

	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,

}

// 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

        %2 = stablehlo.convolution(%1, %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<1x3x3x4xf32>, tensor<3x3x4x4xf32>) -> tensor<1x3x3x4xf32>

  }) {device = "/device:CPU:0"} : () -> (tensor<!tf_type.string>, tensor<2x!tf_type.string>)
  // CHECK-DAG: %[[LAYOUT0:.*]] = "tf.TPUGetLayoutOp"(%[[COMPILE]]#1) <{index = 0 : i64, is_output = false}>
  // CHECK-DAG: %[[LAYOUT1:.*]] = "tf.TPUGetLayoutOp"(%[[COMPILE]]#1) <{index = 1 : i64, is_output = false}>
  // CHECK: %[[ITER:.*]]:2 = "tf.IteratorGetNext"
  %2:2 = "tf.IteratorGetNext"(%arg0) {device = "/device:CPU:0"}

    %0 = stablehlo.convolution(%arg0, %arg1) dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], window = {pad = [[0, 1], [1, 1]]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64} : (tensor<1x3x4x3xf32>, tensor<2x3x3x2xf32>) -> tensor<1x3x4x2xf32>
    return %0 : tensor<1x3x4x2xf32>
  }
  // CHECK: func private @composite_conv_fn
  // CHECK: %[[CONV:.+]] = stablehlo.convolution

    %tensor_int64 = "tf.Const"() {value = dense<1> : tensor<i64>} : () -> tensor<i64>
    %tensor_string = "tf.Const"() {value = dense<""> : tensor<!tf_type.string>} : () -> tensor<!tf_type.string>
    "tf.WriteSummary"(%arg1, %tensor_int64, %tensor_int64, %tensor_string, %tensor_string): (tensor<*x!tf_type.resource>, tensor<i64>, tensor<i64>, tensor<!tf_type.string>, tensor<!tf_type.string>) -> ()