R/159_reductions_dcp2cone_canonicalizers_log_sum_exp_canon.R

In CVXR: Disciplined Convex Optimization

#####
## DO NOT EDIT THIS FILE!! EDIT THE SOURCE INSTEAD: rsrc_tree/reductions/dcp2cone/canonicalizers/log_sum_exp_canon.R
#####

## CVXPY SOURCE: reductions/dcp2cone/canonicalizers/log_sum_exp_canon.py
## log(sum(exp(x))) <= t <=> sum(exp(x - t)) <= 1
## Recursive: calls exp_canon, axis-aware broadcasting of t


log_sum_exp_canon <- function(expr, args, solver_context = NULL) {
  x <- args[[1L]]
  shape <- expr@shape
  axis <- expr@axis
  keepdims <- expr@keepdims
  t <- Variable(shape = shape)

  ## Promote t to match x shape based on axis
  if (is.null(axis)) {
    ## axis=NULL: scalar -> broadcast to x shape
    promoted_t <- cvxr_promote(t, x@shape)
  } else if (axis == 2L) {
    ## axis=2: shape = (1, ncol) -> broadcast rows
    ones_col <- Constant(matrix(1, nrow = x@shape[1L], ncol = 1L))
    t_row <- reshape_expr(t, c(1L, x@shape[2L]))
    promoted_t <- ones_col %*% t_row
  } else {
    ## axis=1: shape = (m, 1) -> broadcast columns
    t_col <- reshape_expr(t, c(x@shape[1L], 1L))
    ones_row <- Constant(matrix(1, nrow = 1L, ncol = x@shape[2L]))
    promoted_t <- t_col %*% ones_row
  }

  exp_expr <- Exp(x - promoted_t)
  exp_result <- exp_canon(exp_expr, exp_expr@args)
  obj <- exp_result[[1L]]
  constraints <- exp_result[[2L]]

  ## Sum along the axis
  obj <- SumEntries(obj, axis = axis, keepdims = keepdims)
  ones <- Constant(matrix(1, nrow = shape[1L], ncol = shape[2L]))
  constraints <- c(constraints, list(obj <= ones))
  list(t, constraints)
}

method(dcp_canonicalize, LogSumExp) <- log_sum_exp_canon
method(has_dcp_canon, LogSumExp) <- function(expr) TRUE