Extends add object to general pickle-like objects

[d3sm0]

cloudpickle doesn't seem to work well with jax. One can use jnp.save for jax tensors but i preferred to be general.