&guaranteed_constant)) {
    return false;
  }
  return guaranteed_constant;
}

// Finds the `index` of an _Arg or _Retval node.
Status GetIndexAttr(const Node& n, int num_args, int* index) {
  TF_RETURN_IF_ERROR(GetNodeAttr(n.attrs(), "index", index));
  if (*index < 0 || *index >= num_args) {
    return errors::InvalidArgument("Invalid ", n.type_string(), " number ",

        for (int i = res.index(), j = 0; i < op->getNumResults(); i++, j++) {
          auto index = builder.create<mlir::arith::ConstantOp>(
              op->getLoc(), builder.getIndexAttr(j));
          auto element_op = builder.create<GetElementOp>(
              op->getLoc(), unconstrainted_tensor_type,
              new_op.getResult(res.index()), index.getResult());

  Value lhs_shape = rewriter->create<shape::ShapeOfOp>(loc, lhs);
  Value rhs_shape = rewriter->create<shape::ShapeOfOp>(loc, rhs);
  Value const_neg2 =
      rewriter->create<arith::ConstantOp>(loc, rewriter->getIndexAttr(-2));
  auto shape_type = shape::ShapeType::get(rewriter->getContext());
  auto lhs_splitted = rewriter->create<shape::SplitAtOp>(
      loc, TypeRange{shape_type, shape_type}, lhs_shape, const_neg2);