How to use the diffcp.solve_and_derivative_batch function in diffcp

for p, sz in zip(params_numpy, ctx.batch_sizes)]
                c, _, neg_A, b = compiler.apply_parameters(
                    dict(zip(param_ids, params_numpy_i)),
                    keep_zeros=True)
                A = -neg_A  # cvxpy canonicalizes -A
                As.append(A)
                bs.append(b)
                cs.append(c)
                cone_dicts.append(cone_dims)
                ctx.shapes.append(A.shape)
            info['canon_time'] = time.time() - start

            # compute solution and derivative function
            start = time.time()
            try:
                xs, _, _, _, ctx.DT_batch = diffcp.solve_and_derivative_batch(
                    As, bs, cs, cone_dicts, **solver_args)
            except diffcp.SolverError as e:
                print(
                    "Please consider re-formulating your problem so that "
                    "it is always solvable or increasing the number of "
                    "solver iterations.")
                raise e
            info['solve_time'] = time.time() - start

            # extract solutions and append along batch dimension
            sol = [[] for _ in range(len(variables))]
            for i in range(ctx.batch_size):
                sltn_dict = compiler.split_solution(
                    xs[i], active_vars=var_dict)
                for j, v in enumerate(variables):
                    sol[j].append(to_torch(

def f_forward():
        return diffcp.solve_and_derivative_batch(As, bs, cs, Ks,
                                                 n_jobs_forward=n_jobs, n_jobs_backward=n_jobs, solver="ECOS", verbose=False)
    xs, ys, ss, D_batch, DT_batch = diffcp.solve_and_derivative_batch(As, bs, cs, Ks,

else:
                    new_id = param.id
                    param_map[new_id] = value
            params = [param_map[pid] for pid in self.param_ids]

        As, bs, cs = [], [], []
        for i in range(batch_size):
            params_i = [
                p if sz == 0 else p[i] for p, sz in zip(params, batch_sizes)]
            A, b, c = self._problem_data_from_params(params_i)
            As.append(A)
            bs.append(b)
            cs.append(c)

        try:
            xs, _, ss, _, DT = diffcp.solve_and_derivative_batch(
                As=As, bs=bs, cs=cs, cone_dicts=[self.cones] * batch_size,
                **solver_args)
        except diffcp.SolverError as e:
            print(
                "Please consider re-formulating your problem so that "
                "it is always solvable or increasing the number of "
                "solver iterations.")
            raise e

        DT = self._restrict_DT_to_dx(DT, batch_size, ss[0].shape)
        solns = [self._split_solution(x) for x in xs]
        # soln[i] is a tensor with first dimension equal to batch_size, holding
        # the optimal values for variable i
        solution = [
            tf.stack([s[i] for s in solns]) for i in range(len(self.vars))]
        if not any_batched:

def time_function(f, N=1):
    result = []
    for i in range(N):
        tic = time.time()
        f()
        toc = time.time()
        result += [toc - tic]
    return np.mean(result), np.std(result)

for n_jobs in range(1, 8):
    def f_forward():
        return diffcp.solve_and_derivative_batch(As, bs, cs, Ks,
                                                 n_jobs_forward=n_jobs, n_jobs_backward=n_jobs, solver="ECOS", verbose=False)
    xs, ys, ss, D_batch, DT_batch = diffcp.solve_and_derivative_batch(As, bs, cs, Ks,
                                                                      n_jobs_forward=1, n_jobs_backward=n_jobs, solver="ECOS", verbose=False)

    def f_backward():
        DT_batch(xs, ys, ss, mode="lsqr")

    mean_forward, std_forward = time_function(f_forward)
    mean_backward, std_backward = time_function(f_backward)
    print("%03d | %4.4f +/- %2.2f | %4.4f +/- %2.2f" %
          (n_jobs, mean_forward, std_forward, mean_backward, std_backward))