op.getLoc(), op.getType(), op.getOperands(),
        builder.getI64ArrayAttr(partition_dims), builder.getBoolAttr(false),
        op.get_XlaShardingAttr());
    op->replaceAllUsesWith(pi);
  } else {
    auto po = builder.create<TF::TPUPartitionedOutputV2Op>(
        op.getLoc(), op.getResultTypes(), op.getOperand(),
        builder.getI64ArrayAttr(partition_dims), op.get_XlaShardingAttr());
    op->replaceAllUsesWith(po);
  }

          conv_op.getLoc(), conv_output_type, sliced_lhs, reshaped_filter,
          rewriter.getI64ArrayAttr(strides),
          /*padding=*/rewriter.getStringAttr(padding),
          /*explicit_paddings=*/rewriter.getI64ArrayAttr(new_padding),
          /*data_format=*/rewriter.getStringAttr("NHWC"),
          /*dilations=*/rewriter.getI64ArrayAttr(dilation));
    } else if (num_spatial_dims == 3) {
      output = rewriter.create<TF::Conv3DOp>(

  const int stride_w = mlir::cast<IntegerAttr>(stride[2]).getInt();
  return rewriter.getI64ArrayAttr({stride_h, stride_w});
}

Attribute GetLhsDilationValue(PatternRewriter& rewriter,
                              llvm::StringMap<Attribute>& identifier_to_attr) {
  return rewriter.getI64ArrayAttr({1, 1});
}

Attribute GetRhsDilationValue(PatternRewriter& rewriter,

  "CreateScalarConstValue<"# type #">($_builder, $_loc, "# value #")">;

// Creates an I64 array attribute with given values.
class CreateI64ArrayAttr<string values> : NativeCodeCall<
  "$_builder.getI64ArrayAttr("# values #")">;

// Creates a string attribute with given values.
class CreateStringAttr<string values> : NativeCodeCall<
  "$_builder.getStringAttr(\""# values #"\")">;

      llvm::ArrayRef<NamedAttribute>{
          builder->getNamedAttr(
              "device_var_reads_indices",
              builder->getI64ArrayAttr(device_var_reads_indices)),
          builder->getNamedAttr(
              "device_var_updates_indices",
              builder->getI64ArrayAttr(device_var_updates_indices))});

  // Wrap in launch for device assignment.
  auto merged_execute_launch = builder->create<tf_device::LaunchOp>(

      dilation_w_factor = 1;
    }
  }

  if (dilation_h_factor == -1 || dilation_w_factor == -1) {
    return {};
  }

  return rewriter.getI64ArrayAttr({1, dilation_h_factor, dilation_w_factor, 1});
}

}  // namespace TFL
}  // namespace mlir

}

static mlir::Attribute BuildI64ArrayAttr(std::vector<int32_t> value,
                                         mlir::Builder builder) {
  std::vector<int64_t> typecast(value.begin(), value.end());
  return builder.getI64ArrayAttr(typecast);
}

static mlir::Attribute BuildVhloBooleanV1Attr(bool value,
                                              mlir::Builder builder) {
  return mlir::vhlo::BooleanV1Attr::get(builder.getContext(), value);

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

      replicated_inputs, packed_inputs, result_types);

  if (!mirrored_variable_indices.empty())
    replicate_op->setAttr(kMirroredVariableIndicesAttr,
                          builder.getI64ArrayAttr(mirrored_variable_indices));

  // Replace replicated cluster results with replicate op results.
  uint64_t offset = 0;
  for (auto [idx, result] : llvm::enumerate(cluster.getResults())) {

}

TEST(TPURewriteDeviceUtilTest, TestGetDeviceCoordinates) {
  mlir::MLIRContext context;
  mlir::Builder builder(&context);
  auto device_assignment_attr = builder.getI64ArrayAttr({1, 2, 3});
  auto status_or_device_coodinates =
      GetDeviceCoordinates(device_assignment_attr);
  ASSERT_TRUE(status_or_device_coodinates.ok());
  auto device_coordinates = status_or_device_coodinates.value();