// CHECK-NEXT: tfr.return %[[bl]] : !tfr.tensor_list
}

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

  // CHECK: tac.cost = 0x4D5C0000
  %0 = "tfl.conv_2d"(%arg0, %arg1, %arg2) {dilation_h_factor = 1 : i32, dilation_w_factor = 1 : i32, padding = "SAME", stride_h = 1 : i32, stride_w = 1 : i32, fused_activation_function = "RELU6", tac.device = "CPU"} : (tensor<256x32x32x3xf32>, tensor<16x3x3x3xf32>, tensor<16xf32>) -> tensor<256x32x32x16xf32>
  func.return %0 : tensor<256x32x32x16xf32>
}

      return failure();
    }

    // Constructs strides array.
    // For example, [2, 3] -> [1, 2, 3, 1].
    SmallVector<int64_t, num_spatial_dims + 2> strides({1});
    for (const auto v :
         conv_op.getWindowStrides().value().getValues<int64_t>()) {
      strides.emplace_back(v);
    }
    strides.emplace_back(1);

    // Constructs dilation array.

    for (Striped<?> striped : allImplementations()) {
      assertBasicInvariants(striped);
    }
  }

  private static void assertBasicInvariants(Striped<?> striped) {
    Set<Object> observed = Sets.newIdentityHashSet(); // for the sake of weakly referenced locks.
    // this gets the stripes with #getAt(index)
    for (int i = 0; i < striped.size(); i++) {
      Object object = striped.getAt(i);
      assertNotNull(object);

    for (Striped<?> striped : allImplementations()) {
      assertBasicInvariants(striped);
    }
  }

  private static void assertBasicInvariants(Striped<?> striped) {
    Set<Object> observed = Sets.newIdentityHashSet(); // for the sake of weakly referenced locks.
    // this gets the stripes with #getAt(index)
    for (int i = 0; i < striped.size(); i++) {
      Object object = striped.getAt(i);
      assertNotNull(object);

  // CHECK-NEXT:       outputs: [ 1 ],
  // CHECK-NEXT:       builtin_options_type: Pool2DOptions,
  // CHECK-NEXT:       builtin_options: {
  // CHECK-NEXT:         padding: VALID,
  // CHECK-NEXT:         stride_w: 1,
  // CHECK-NEXT:         stride_h: 3,
  // CHECK-NEXT:         filter_width: 6,
  // CHECK-NEXT:         filter_height: 3
  // CHECK-NEXT:       }
  // CHECK-NEXT:     } ]
  // CHECK-NEXT:    name: "main"
  // CHECK-NEXT:   } ],

  %3 = "tfl.conv_2d"(%0, %1, %2) {dilation_h_factor = 1 : i32, dilation_w_factor = 1 : i32, fused_activation_function = "NONE", padding = "SAME", stride_h = 2 : i32, stride_w = 2 : i32} : (tensor<1x224x224x3x!quant.uniform<i8:f32, 7.812500e-03>>, tensor<32x3x3x3x!quant.uniform<i8<-127:127>:f32, 0.021826678373682216>>, tensor<32x!quant.uniform<i32:f32, 1.7052092479439231E-4>>) -> tensor<1x112x112x32x!quant.uniform<i8:f32,...

  GetI32VectorFromDenseI64CompositeAttr(composite_attrs, "padding", &padding);
  data.ph = padding[0];
  data.pw = padding[1];

  std::vector<int32_t> strides;
  GetI32VectorFromDenseI64CompositeAttr(composite_attrs, "stride", &strides);
  data.sh = strides[0];
  data.sw = strides[1];

  data.ceil_mode =
      GetBoolFromCompositeAttr(composite_attrs, "ceil_mode").value();

  if (strides[0] != 1 || strides[3] != 1) return 1;

  // Space to depth only supports height_stride == width_stride case.
  if (strides[1] != strides[2]) return 1;

  return strides[1];
}

void TPUSpaceToDepthPass::runOnOperation() {
  std::optional<tf_device::ClusterFuncOp> cluster_func;
  // Space to depth only supports training loop.

  (TFL_TransposeConvOp
    $output_shape,
    (TFL_DequantizeOp $quant_weights),
    $quant_input,
    $bias, $padding, $stride_h, $stride_w, $faf),
  (TFL_TransposeConvOp $output_shape, $quant_weights,
    $quant_input, $bias, $padding, $stride_h, $stride_w, $faf),