cloudpickle fails only when jitting after de-serialisation

[simon-bachhuber]

Not an issue but a question. Why does cloudpickle work partially in this example?

import equinox as eqx
import jax.numpy as jnp 
import cloudpickle as p 

class Model(eqx.Module):
    x: jnp.ndarray
    def __call__(self):
        return self.x

m = Model(jnp.ones((3,3)))
m = p.loads(p.dumps(m))

# this works
m()
# this does not
eqx.filter_jit(m)()
>> AttributeError: 'Model' object has no attribute 'x'

[patrick-kidger]

Looks like some kind of bug in cloudpickle, mutating what counts as the fields of a dataclass. Here's a reproducer that doesn't use Equinox or JAX:

import cloudpickle as p
import dataclasses

@dataclasses.dataclass
class Model:
    x: int

print(dataclasses.fields(Model))
p.loads(p.dumps(Model(1)))
print(dataclasses.fields(Model))

[simon-bachhuber]

Ha. I will move this then over to cloudpickle directly.