Calling jax.jacobian results in error

[GeoffNN]

Hi,
Thanks for the cool package!
I've just started playing around with it on a QP I need to differentiate through. I'm getting the following error when calling jax.jacobian on the layer. The forward pass works fine, and finds a nontrivial solution as indicated by the verbose output of the solver.
Any help here would be appreciated!

Solver output:

----------------------------------------------------------------------------
	SCS v2.1.4 - Splitting Conic Solver
	(c) Brendan O'Donoghue, Stanford University, 2012
----------------------------------------------------------------------------
Lin-sys: sparse-direct, nnz in A = 20101
eps = 1.00e-05, alpha = 1.50, max_iters = 5000, normalize = 1, scale = 1.00
acceleration_lookback = 10, rho_x = 1.00e-03
Variables n = 102, constraints m = 202
Cones:	primal zero / dual free vars: 100
	soc vars: 102, soc blks: 1
Setup time: 4.14e-02s
----------------------------------------------------------------------------
 Iter | pri res | dua res | rel gap | pri obj | dua obj | kap/tau | time (s)
----------------------------------------------------------------------------
     0| 2.91e+19  2.29e+19  1.00e+00 -6.29e+19  1.45e+18  5.11e+19  2.14e-02 
   100| 3.48e-05  1.08e-05  5.48e-05  5.81e-02  5.82e-02  3.46e-18  4.10e-02 
   200| 3.62e-03  1.26e-03  1.62e-03  4.83e-02  4.65e-02  2.10e-17  5.60e-02 
   300| 1.37e-01  1.28e-02  3.03e-02  1.30e-02  4.51e-02  2.48e-17  6.72e-02 
   400| 5.11e-02  1.60e-02  2.01e-02  7.96e-02  5.67e-02  2.13e-17  7.94e-02 
   500| 1.51e-02  5.96e-03  5.18e-03  4.17e-02  3.61e-02  7.22e-18  9.34e-02 
   600| 4.03e-04  1.67e-04  1.32e-04  3.91e-02  3.92e-02  3.66e-17  1.05e-01 
   700| 1.25e-03  2.41e-04  1.02e-03  3.82e-02  3.93e-02  1.07e-18  1.18e-01 
   800| 8.28e-03  2.12e-03  2.41e-03  3.25e-02  3.51e-02  2.42e-17  1.30e-01 
   900| 9.51e-04  4.79e-04  1.53e-04  3.87e-02  3.88e-02  3.10e-17  1.42e-01 
  1000| 1.73e-03  1.01e-03  1.14e-03  3.82e-02  3.95e-02  7.56e-18  1.58e-01 
  1100| 2.17e-03  1.14e-03  5.81e-04  3.79e-02  3.86e-02  6.94e-18  1.76e-01 
  1200| 2.57e-02  9.40e-03  5.71e-03  3.29e-02  3.91e-02  4.66e-17  1.93e-01 
  1300| 5.11e-03  2.51e-03  7.59e-03  4.72e-02  3.90e-02  5.41e-18  2.07e-01 
  1400| 1.46e-03  3.23e-04  5.90e-06  3.68e-02  3.68e-02  3.53e-18  2.19e-01 
  1500| 2.58e-03  2.64e-04  1.59e-03  3.64e-02  3.81e-02  5.11e-18  2.36e-01 
  1600| 1.07e-03  5.72e-04  2.59e-04  3.77e-02  3.80e-02  1.17e-17  2.49e-01 
  1700| 1.21e-03  5.51e-04  9.79e-04  3.82e-02  3.72e-02  2.62e-17  2.63e-01 
  1800| 6.70e-04  1.30e-04  2.71e-04  3.71e-02  3.74e-02  3.90e-17  2.76e-01 
  1900| 2.35e-02  1.15e-02  4.02e-03  4.06e-02  3.62e-02  3.50e-17  2.88e-01 
  2000| 1.21e-03  1.82e-04  1.99e-03  3.83e-02  3.61e-02  1.90e-17  3.01e-01 
  2100| 4.19e-02  1.53e-02  1.78e-02  5.94e-02  3.98e-02  4.11e-19  3.12e-01 
  2200| 9.61e-04  3.36e-04  5.88e-04  3.86e-02  3.79e-02  1.78e-17  3.27e-01 
  2300| 1.10e-03  5.49e-04  7.74e-04  3.86e-02  3.78e-02  2.92e-17  3.40e-01 
  2400| 5.45e-03  1.52e-03  2.36e-03  4.30e-02  4.04e-02  7.82e-17  3.57e-01 
  2500| 1.66e-02  9.82e-03  1.03e-02  5.18e-02  4.06e-02  3.10e-17  3.69e-01 
  2600| 1.31e-04  7.45e-05  9.50e-05  3.75e-02  3.76e-02  1.97e-17  3.86e-01 
  2700| 4.12e-04  2.23e-04  4.29e-04  3.79e-02  3.75e-02  5.47e-18  3.99e-01 
  2800| 2.99e-04  9.80e-05  3.23e-05  3.76e-02  3.76e-02  1.43e-17  4.11e-01 
  2900| 1.62e-03  9.50e-04  2.36e-03  3.44e-02  3.69e-02  2.24e-17  4.23e-01 
  3000| 5.68e-04  3.09e-04  3.79e-04  3.79e-02  3.75e-02  4.01e-18  4.35e-01 
  3100| 1.02e-02  5.00e-03  1.26e-03  4.00e-02  4.14e-02  2.91e-17  4.47e-01 
  3200| 9.33e-03  4.59e-03  3.19e-03  4.10e-02  3.75e-02  6.83e-18  4.60e-01 
  3300| 1.20e-03  6.14e-04  7.47e-04  3.67e-02  3.75e-02  9.50e-18  4.72e-01 
  3400| 4.41e-02  2.47e-02  2.23e-02  1.80e-02  4.17e-02  2.27e-17  4.84e-01 
  3500| 6.01e-03  3.32e-03  2.26e-03  2.76e-02  2.52e-02  5.72e-17  4.96e-01 
  3600| 7.45e-04  8.20e-05  9.78e-05  3.70e-02  3.69e-02  1.19e-17  5.08e-01 
  3700| 1.16e-04  4.97e-05  1.93e-04  3.72e-02  3.70e-02  2.38e-17  5.20e-01 
  3800| 3.30e-03  1.48e-03  5.65e-04  3.74e-02  3.67e-02  6.22e-18  5.32e-01 
  3900| 7.25e-03  3.55e-03  1.21e-03  3.47e-02  3.35e-02  3.64e-18  5.44e-01 
  4000| 1.61e-03  3.41e-04  4.68e-04  3.70e-02  3.65e-02  2.64e-17  5.60e-01 
  4100| 5.24e-02  1.16e-02  1.15e-02  6.02e-02  4.74e-02  6.24e-18  5.73e-01 
  4200| 4.50e-03  2.61e-03  1.55e-03  4.08e-02  3.91e-02  1.08e-18  5.84e-01 
  4300| 7.94e-03  4.51e-03  8.34e-04  3.66e-02  3.75e-02  3.63e-17  6.00e-01 
  4400| 3.09e-03  1.80e-03  2.32e-03  3.97e-02  4.22e-02  3.98e-17  6.12e-01 
  4500| 6.47e-04  2.21e-04  1.00e-03  3.60e-02  3.70e-02  6.48e-18  6.24e-01 
  4600| 1.02e-02  4.87e-03  7.15e-03  4.29e-02  3.51e-02  1.17e-17  6.35e-01 
  4700| 6.71e-04  3.42e-04  2.42e-04  3.68e-02  3.71e-02  1.07e-18  6.48e-01 
  4800| 2.58e-03  1.50e-03  2.71e-04  3.53e-02  3.56e-02  1.59e-17  6.59e-01 
  4900| 2.89e-01  1.71e-01  7.00e-02 -1.56e-01 -6.99e-02  1.67e-17  6.71e-01 
  5000| 6.05e-03  2.44e-03  7.90e-04  3.51e-02  3.59e-02  1.61e-17  6.83e-01 
----------------------------------------------------------------------------
Status: Solved/Inaccurate
Hit max_iters, solution may be inaccurate, returning best found solution.
Timing: Solve time: 6.83e-01s
	Lin-sys: nnz in L factor: 25455, avg solve time: 7.81e-05s
	Cones: avg projection time: 3.44e-07s
	Acceleration: avg step time: 3.39e-05s
----------------------------------------------------------------------------
Error metrics:
dist(s, K) = 5.0672e-20, dist(y, K*) = 0.0000e+00, s'y/|s||y| = 3.5284e-16
primal res: |Ax + s - b|_2 / (1 + |b|_2) = 1.8102e-05
dual res:   |A'y + c|_2 / (1 + |c|_2) = 4.0469e-06
rel gap:    |c'x + b'y| / (1 + |c'x| + |b'y|) = 9.8806e-06
----------------------------------------------------------------------------
c'x = 0.0391, -b'y = 0.0391
============================================================================
----------------------------------------------------------------------------
	SCS v2.1.4 - Splitting Conic Solver
	(c) Brendan O'Donoghue, Stanford University, 2012
----------------------------------------------------------------------------
Lin-sys: sparse-direct, nnz in A = 20101
eps = 1.00e-05, alpha = 1.50, max_iters = 5000, normalize = 1, scale = 1.00
acceleration_lookback = 0, rho_x = 1.00e-03
Variables n = 102, constraints m = 202
Cones:	primal zero / dual free vars: 100
	soc vars: 102, soc blks: 1
Setup time: 6.42e-02s
----------------------------------------------------------------------------
 Iter | pri res | dua res | rel gap | pri obj | dua obj | kap/tau | time (s)
----------------------------------------------------------------------------
     0| 2.91e+19  2.29e+19  1.00e+00 -6.29e+19  1.45e+18  5.11e+19  5.68e-02 
   100| 3.87e-05  2.47e-07  5.24e-09  5.82e-02  5.82e-02  2.21e-17  7.08e-02 
   140| 6.52e-06  1.66e-07  5.32e-09  5.82e-02  5.82e-02  2.23e-17  7.57e-02 
----------------------------------------------------------------------------
Status: Solved
Timing: Solve time: 7.57e-02s
	Lin-sys: nnz in L factor: 25455, avg solve time: 1.01e-04s
	Cones: avg projection time: 7.41e-07s
	Acceleration: avg step time: 8.07e-08s
----------------------------------------------------------------------------
Error metrics:
dist(s, K) = 4.4318e-19, dist(y, K*) = 1.1102e-16, s'y/|s||y| = 3.7957e-17
primal res: |Ax + s - b|_2 / (1 + |b|_2) = 6.5182e-06
dual res:   |A'y + c|_2 / (1 + |c|_2) = 1.6561e-07
rel gap:    |c'x + b'y| / (1 + |c'x| + |b'y|) = 5.3244e-09
----------------------------------------------------------------------------
c'x = 0.0582, -b'y = 0.0582
============================================================================

Error:

/usr/local/lib/python3.7/dist-packages/jax/_src/api.py in jacfun(*args, **kwargs)
    970     tree_map(partial(_check_input_dtype_jacfwd, holomorphic), dyn_args)
    971     pushfwd = partial(_jvp, f_partial, dyn_args)
--> 972     y, jac = vmap(pushfwd, out_axes=(None, -1))(_std_basis(dyn_args))
    973     tree_map(partial(_check_output_dtype_jacfwd, holomorphic), y)
    974     example_args = dyn_args[0] if isinstance(argnums, int) else dyn_args

<ipython-input-60-2f66503cf065> in <lambda>(*args, **kwargs)
     16   return lambda *args, **kwargs: CvxpyLayer(prob,
     17                     parameters=[_F, _dx_F, _dt_F, _ground_truth],
---> 18                     variables=[_lam])(*args, **kwargs)[0]

/usr/local/lib/python3.7/dist-packages/cvxpylayers/jax/cvxpylayer.py in CvxpyLayer(problem, parameters, variables, gp)
     74         )
     75     else:
---> 76         data, _, _ = problem.get_problem_data(solver=cp.SCS)
     77         compiler = data[cp.settings.PARAM_PROB]
     78         param_ids = [p.id for p in param_order]

/usr/local/lib/python3.7/dist-packages/cvxpy/problems/problem.py in get_problem_data(self, solver, gp, enforce_dpp, verbose)
    599 
    600             data, solver_inverse_data = solving_chain.solver.apply(
--> 601                 self._cache.param_prog)
    602             inverse_data = self._cache.inverse_data + [solver_inverse_data]
    603             self._compilation_time = time.time() - start

/usr/local/lib/python3.7/dist-packages/cvxpy/reductions/solvers/conic_solvers/scs_conif.py in apply(self, problem)
    219         # Apply parameter values.
    220         # Obtain A, b such that Ax + s = b, s \in cones.
--> 221         c, d, A, b = problem.apply_parameters()
    222         data[s.C] = c
    223         inv_data[s.OFFSET] = d

/usr/local/lib/python3.7/dist-packages/cvxpy/reductions/dcp2cone/cone_matrix_stuffing.py in apply_parameters(self, id_to_param_value, zero_offset, keep_zeros)
    203             self.reduced_A, param_vec, self.x.size,
    204             nonzero_rows=self._A_mapping_nonzero, with_offset=True,
--> 205             problem_data_index=self.problem_data_index)
    206         return c, d, A, np.atleast_1d(b)
    207 

/usr/local/lib/python3.7/dist-packages/cvxpy/cvxcore/python/canonInterface.py in get_matrix_from_tensor(problem_data_tensor, param_vec, var_length, nonzero_rows, with_offset, problem_data_index)
    236     if nonzero_rows is not None and nonzero_rows.size > 0:
    237         A_nrows, _ = A.shape
--> 238         A_rows, A_cols = nonzero_csc_matrix(A)
    239         A_vals = np.append(A.data, np.zeros(nonzero_rows.size))
    240         A_rows = np.append(A_rows, nonzero_rows % A_nrows)

/usr/local/lib/python3.7/dist-packages/cvxpy/cvxcore/python/canonInterface.py in nonzero_csc_matrix(A)
    144     # this function returns (rows, cols) corresponding to nonzero entries in
    145     # A; an entry that is explicitly set to zero is treated as nonzero
--> 146     assert not np.isnan(A.data).any()
    147 
    148     # scipy drops rows, cols with explicit zeros; use nan as a sentinel

AssertionError: