TypeError: len() of unsized object when trying to solve with dccp

[njwfish]

Issue

When I try to solve my problem I get the following error:

Traceback (most recent call last):
  File "/Users/njwfish/miniforge3/envs/ei/lib/python3.8/site-packages/IPython/core/interactiveshell.py", line 3457, in run_code
    exec(code_obj, self.user_global_ns, self.user_ns)
  File "<ipython-input-129-984a6a5f0728>", line 1, in <module>
    prob.solve(solver="ECOS", method="dccp")
  File "/Users/njwfish/miniforge3/envs/ei/lib/python3.8/site-packages/cvxpy/problems/problem.py", line 481, in solve
    return solve_func(self, *args, **kwargs)
  File "/Users/njwfish/miniforge3/envs/ei/lib/python3.8/site-packages/dccp/problem.py", line 57, in dccp
    result_temp = iter_dccp(
  File "/Users/njwfish/miniforge3/envs/ei/lib/python3.8/site-packages/dccp/problem.py", line 240, in iter_dccp
    temp = convexify_constr(arg)
  File "/Users/njwfish/miniforge3/envs/ei/lib/python3.8/site-packages/dccp/constraint.py", line 76, in convexify_constr
    right = linearize(constr.args[1])
  File "/Users/njwfish/miniforge3/envs/ei/lib/python3.8/site-packages/dccp/linearize.py", line 58, in linearize
    grad_map = expr.grad
  File "/Users/njwfish/miniforge3/envs/ei/lib/python3.8/site-packages/cvxpy/atoms/atom.py", line 402, in grad
    grad_self = self._grad(arg_values)
  File "/Users/njwfish/miniforge3/envs/ei/lib/python3.8/site-packages/cvxpy/atoms/log_sum_exp.py", line 51, in _grad
    return self._axis_grad(values)
  File "/Users/njwfish/miniforge3/envs/ei/lib/python3.8/site-packages/cvxpy/atoms/axis_atom.py", line 106, in _axis_grad
    D = D + sp.csc_matrix((np.array(d)[0], (row, col)),
  File "/Users/njwfish/miniforge3/envs/ei/lib/python3.8/site-packages/scipy/sparse/compressed.py", line 54, in __init__
    other = self.__class__(coo_matrix(arg1, shape=shape,
  File "/Users/njwfish/miniforge3/envs/ei/lib/python3.8/site-packages/scipy/sparse/coo.py", line 196, in __init__
    self._check()
  File "/Users/njwfish/miniforge3/envs/ei/lib/python3.8/site-packages/scipy/sparse/coo.py", line 281, in _check
    if self.nnz > 0:
  File "/Users/njwfish/miniforge3/envs/ei/lib/python3.8/site-packages/scipy/sparse/base.py", line 246, in nnz
    return self.getnnz()
  File "/Users/njwfish/miniforge3/envs/ei/lib/python3.8/site-packages/scipy/sparse/coo.py", line 241, in getnnz
    nnz = len(self.data)
TypeError: len() of unsized object

I think this problem may be related to the indexing I am doing to solve this problem, but I am not sure. I have verified my problem is dccp. The issue persists across different solvers. Additionally if I drop the equality constraint and make the problem dcp then the problem is solved as I've formulated it both using prob.solve(solver="ECOS") and prob.solve(solver="ECOS", method="dccp").

So that's the issue, I'll give some background on the problem now so the example makes sense. I'm sure there's a more minimal replicating example but this is my use case. Any help would be much appreciated!

Background and code

I am trying to solve a difference of logsumexps problem in rational choice theory. Essentially you have a set of people who make choose the best and worst item from a set. Both the items and the people are parameterized by feature vectors. You model the choice as a linear combination of these two feature vectors pushed through a softmax over the elements presented in the choice set.

This is a little tricky to formulate in cvxpy because it the choice sets are [N, C, K] tensors, the respondents [N, D] matrices, and the parameter we are optimizing is a [K, D] matrix. We want to take a matrix vector product only along the last dimension of the choice set, leaving us with a [N, C] matrix of the resulting logits. This can be achieved using an einsum logits = np.einsum('nck,kd,nd->nc', choice_set_features, W, respondent_features).

To illustrate:

import numpy as np
import cvxpy as cp

# jax for a nice categorical sampler
import jax

# generate data

N = 1000
K = 3
C = 3
D = 3

# simulated covariates
respondent_features = np.random.normal(size=(N, D))
choice_set_features = np.random.normal(size=(N, C, K))

# parameter we're estimating
W = np.random.normal(size=(K, D))

logits = np.einsum('nck,kd,nd->nc', choice_set_features, W, respondent_features)

# sample the best
best = jax.random.categorical(jax.random.PRNGKey(1), logits)

# sample the worst by sampling from 1 - p_ij
expits = np.exp(logits)
probs = expits / np.sum(expits, axis=1)[:, None]
compliment_logits = np.log(1 - probs)
worst = jax.random.categorical(jax.random.PRNGKey(1), compliment_logits)

# set up cvx parameters
t = cp.Variable(N)
hatW = cp.Variable((K, D))

# compute logits with cvx parameters, this bit is
# replicating the einsum with cvx native code
def vectorized_dot_product(a, b):
    return cp.sum(cp.multiply(a, b), axis=1)

W_respondent_features = respondent_features @ hatW.T
logits = cp.vstack([
    vectorized_dot_product(choice_set_features[:, i, :], W_respondent_features)
    for i in range(C)
]).T

# get best logits with indexing
best_logit = logits[np.arange(N), best]

# get all indices which are not the worst element using indexing
idx_not_worst = np.where(np.tile(np.arange(C)[None, :], (N, 1)) != worst[:, None])
idx_not_worst = idx_not_worst[0].reshape(-1, C - 1), idx_not_worst[1].reshape(-1, C - 1)
not_worst_logits = logits[idx_not_worst]

# set up constraint for dccp form of difference of logsumexps problem
const = t == cp.log_sum_exp(not_worst_logits, axis=1)
nll = cp.sum(cp.multiply(2, cp.log_sum_exp(logits, axis=1)) - best_logit - t)

We can then attempt to solve:

prob = cp.Problem(cp.Minimize(nll), constraints=[const])
prob.solve(solver="ECOS", method="dccp")

resulting in the above error it is easy to check the problem is in fact dccp. If we modify the problem to be dcp instead we can solve with both standard cvx and dccp:

nll = cp.sum(cp.multiply(2, cp.log_sum_exp(logits, axis=1)) - best_logit)
prob = cp.Problem(cp.Minimize(nll))
prob.solve(solver="ECOS")
prob.solve(solver="ECOS", method="dccp")