input.getType().cast<ShapedType>().getElementType());
    inferredReturnTypes.assign({result_type});
    return success();
  }

  auto stride_h = op.getStrideHAttr().getInt();
  auto stride_w = op.getStrideWAttr().getInt();
  auto dilation_h = op.getDilationHFactorAttr().getInt();
  auto dilation_w = op.getDilationWFactorAttr().getInt();

      return op.emitOpError() << "requires begin, end and strides to have the "
                                 "same number of elements";
    }
  }

  // If strides are constants, verify that none of the element is zero.
  DenseIntElementsAttr strides;
  if (matchPattern(op.getStrides(), m_Constant(&strides))) {
    if (llvm::is_contained(strides.getValues<APInt>(), 0))

kernel\_size, 3, out\_channel]. Space to depth is to transform this first
convolution's input to [batch\_size, height // stride, width // stride, 3 \*
stride \* stride] and the kernel to [kernel\_size // stride, kernel\_size //
stride, 3 \* stride \* stride, out\_channel] to improve TPU MXU utilization.

![drawings](images/space_to_depth_transform.png)

      assert(index[i] >= 0);
    }
#endif  // NDEBUG
    int64_t offset = 0;
    int64_t stride = 1;
    for (int64_t dim = shape.size() - 1; dim >= 0; --dim) {
      offset += index[dim] * stride;
      stride *= shape[dim];
    }
    return offset;
  }

  int64_t offset_;
  int64_t stride_;
  int64_t size_;
};

template <typename T>

  dltensor_in->dtype = {kDLFloat, 32, 1};
  dltensor_in->shape = shape.data();
  dltensor_in->strides = strides.data();
  TFE_TensorHandle* handle = TFE_HandleFromDLPack(&dlm_in, status, ctx);
  ASSERT_NE(handle, nullptr)
      << TF_Message(status) << " (shape=[" << absl::StrJoin(shape, ",")
      << "], strides=[" << absl::StrJoin(strides, ",") << "])";

  auto* dlm_out =

  // CHECK: %[[CST_0:.*]] = "tf.Const"() <{value = dense<[0, 3, 1, 2]> : tensor<4xi64>}>
  // CHECK: "tf.Transpose"(%[[R1]], %[[CST_0]])
  %0 = "tf.MaxPool"(%arg0)
       {
         data_format = "NCHW",
         ksize = [1, 1, 3, 3],
         padding = "SAME",
         strides = [1, 1, 2, 2]

  // 1) nullptr indicates tensor is compact and row-majored.
  // 2) fill in the strides array as the real case for compact row-major data.
  // Here we choose option 2, since some frameworks didn't handle the strides
  // argument properly.
  dlm_tensor->dl_tensor.strides = stride_arr->data();

  dlm_tensor->dl_tensor.byte_offset =
      0;  // TF doesn't handle the strides and byte_offsets here
  return static_cast<void*>(dlm_tensor);
}

  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.

  // CHECK-SAME: data_format = "NCHW"
  // CHECK-SAME: dilations = [1, 4, 2, 3]
  // CHECK-SAME: explicit_paddings = [1, 2, 7, 8, 3, 4, 5, 6]
  // CHECK-SAME: padding = "EXPLICIT"
  // CHECK-SAME: strides = [5, 8, 6, 7]
  // CHECK-SAME: (tensor<1x3x32x32xf32>, tensor<1x1x3x8xf32>) -> tensor<1x8x7x7xf32>

  // CHECK: %[[RES_PERM:.*]] = "tf.Const"() <{value = dense<[0, 2, 3, 1]> : tensor<4xi64>}>