Cvxpylayers with lax.scan

[samku]

Hello all,

I am trying to use the differentiable optimizer functionality of cvxpy layers in my code. Since I have a looping operation, I use lax.scan for efficient looping. However, I face the following issue, illustrated through a simple example:

The optimization problem is defined here

a = jnp.array([1.])
def create_opt_layer():
    c = cp.Parameter(1)
    z = cp.Variable(1)
    objective = cp.Minimize(c*z)
    constraints = [z >= 0]
    problem = cp.Problem(objective, constraints)
    qp_layer = CvxpyLayer(problem, parameters=[c], variables=[z])
    return qp_layer

opt_layer = create_opt_layer()

When I call the problem, it works fine
x = jnp.array([2.])
cost_vec= a*x
p_opt, = opt_layer(cost_vec)
print('This works:', jnp.array(p_opt)) 

However, if I want to call the problem in a lax.scan loop as follows:
def solve_LPs(z, x):
    #z is dummy carry
    cost_vec = a*x
    p_opt, = opt_layer(cost_vec)
    p_opt = jnp.array(p_opt)
    return z, p_opt

x = jnp.array([[3.],[4.],[5.]])
f = partial(solve_LPs)
_, p_next = jax.lax.scan(f, 0, x)

It says

line 89, in CvxpyLayerFn
    return CvxpyLayerFn_p.bind(solver_args, *params)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
TypeError: Value {} with type <class 'dict'> is not a valid JAX type

I am lost on how to try resolving it.
Thank you!

[dfm]

It would probably make sense to open this as an issue on the https://github.com/cvxgrp/cvxpylayers repo rather than here since that issue is deep in that code, rather than directly a JAX issue.

[samku]

Thank you! I will ask the question there.