ElementsAttr input, int column) {
  auto int_attr = mlir::cast<DenseIntElementsAttr>(input);
  auto shaped_type = int_attr.getType();
  auto shape = shaped_type.getShape();

  if (shape.size() != 2) return DenseIntElementsAttr();

  llvm::SmallVector<int64_t, 4> values;
  values.reserve(shaped_type.getNumElements() / shape[1]);

  for (const auto &it : llvm::enumerate(int_attr.getValues<APInt>())) {