// TODO(wangtao): add a pass to check if it is profitable to space to depth
// transform and invoke the transform if it is needed.
struct TPUSpaceToDepthPass
    : public impl::TPUSpaceToDepthPassBase<TPUSpaceToDepthPass> {
  void runOnOperation() override;
};

// Updates func argument type to have the updated input shape.
void UpdateFuncType(func::FuncOp func) {

    that are contained in a tf_device.cluster op.
  }];

  let constructor = "TFTPU::CreateTPUClusterCleanupAttributesPass()";
}

def TPUSpaceToDepthPass : Pass<"tf-tpu-space-to-depth-pass", "ModuleOp"> {
  let summary = "Applies automatic space to depth transform for the first or frontier convolutions consume host inputs on TPU.";

  let description = [{