Cells are no longer pickle-able after ba9772b

[bantin]

This commit introduced a closure which is stored on the cell. Now, if I try to pickle a cell I get this:

AttributeError: Can't get local object '_make_dhs_solve.._solve'

Is it reasonable to expect to be able to pickle a cell? It's certainly useful for a variety of scenarios. One workaround is to delete that attribute from the cell object before pickling, but it's a bit annoying to do that every time.

[jnsbck]

Hey,
they should be picklable and we test this here, somehow this slipped past the tests. Do you have a minimal example? :) That would be super helpful to find out what causes this and fix it.

EDIT: I had a brief look, but I cannot find the object _make_dhs_solve._solve that causes this anywhere in the codebase.

Best,

Jonas

[bantin]

I think it is only triggered if you call cell.integrate before pickling. The test doesn't call jx.integrate as far as I can tell. This reproduces the error for me with the latest jaxley

import pickle
import jaxley as jx

comp = jx.Compartment()
branch = jx.Branch(comp, ncomp=4)
cell = jx.Cell(branch, parents=[-1])
cell.record("v")

jx.integrate(cell, t_max=1.0, delta_t=0.025)

print("Attempting pickle.dump(cell) ...")

with open("/tmp/blahblah", "wb") as f:
    pickle.dump(cell, f)

This gives:

AttributeError: Can't get local object '_make_dhs_solve.<locals>._solve'

[bantin]

I believe the troublesome line is in solver_voltage.py:94, though I could be wrong.

    @jax.custom_jvp
    def _solve(diags, lowers, uppers, solves):
        return _raw_solve(diags, lowers, uppers, solves)

[jnsbck]

The offending attribute is cell._dhs_solve_indexer.pop[('_cached_solve_fn', False, cell.ncomp)], which stores the compiled solve function directly on the solve_indexer. As you said, deleting it cell._dhs_solve_indexer.pop(('_cached_solve_fn', False, cell.ncomp)) should safely work since, upon running another jx.integrate on the unpickled object, the solve_indexer gets reinitialized.

The troublesome line is:

jaxley/jaxley/solver_voltage.py

Line 139 in ba9772b

solve_indexer[cache_key] = solve_fn

@michaeldeistler do we want to support pickling after integration? Then I could add something like jx.pickle or module.prepare_pickle() that deletes offending attrs. before pickling