}
    auto in_manual = builder.create<mlir::TF::XlaSpmdFullToShardShapeOp>(
        loc, shard_type, in, common_split_sharding, /*dim=*/-1,
        /*unspecified_dims=*/builder.getI64ArrayAttr({}));
    manual_inputs.push_back(in_manual);
  }

  // Create the _XlaHostComputeMlirOp
  auto host_compute = builder.create<mlir::TF::_XlaHostComputeMlirOp>(
      loc, shard_output_types, manual_inputs,

        rfft_op.getLoc(), rfft2d_out_type, expanded_input.getResult(),
        expanded_fft_len.getResult());

    // Insert the squeeze op.
    auto squeeze_dim = rewriter.getI64ArrayAttr({-2});
    TF::SqueezeOp squeeze = rewriter.create<TF::SqueezeOp>(
        rfft_op.getLoc(), output_type, rfft2d.getResult(), squeeze_dim);

    rewriter.replaceOp(op, squeeze.getResult());