Test pickle support

[jheek]

PR #1475 fixes cloudpickle support but doesn't support add any testing.

Our current test setup seems to import dill indirectly which is incompatible with cloudpickle. Ideally we want to setup testing such that both are tested but don't get in the way of each other. Also we will need to think about how we can test that internal Flax state isn't serialized like was previously the case.

[mattbahr]

It looks like a couple tests have been added for cloudpickle support in linen. Should this issue be closed?

flax/tests/linen/linen_module_test.py

Lines 2558 to 2592 in a9ef59b

	def test_cloudpickle_class(self):
	import cloudpickle
	class MyModule(nn.Module):
	pass
	a = MyModule()
	UnpickledMyModule = cloudpickle.loads(cloudpickle.dumps(MyModule))
	b = UnpickledMyModule()
	def test_cloudpickle_module(self):
	from cloudpickle import cloudpickle_fast
	class NNModuleWithProperty(nn.Module):
	a: int
	b: str
	@property
	def my_property(self):
	return self.b * self.a
	m = NNModuleWithProperty(a=2, b='ok')
	with TemporaryDirectory() as tmpdir:
	filename = f'{tmpdir}/module.pkl'
	with open(filename, 'wb') as f:
	cloudpickle_fast.dump(m, f)
	with open(filename, 'rb') as f:
	obj_loaded = cloudpickle_fast.load(f)
	self.assertEqual(obj_loaded.a, 2)
	self.assertEqual(obj_loaded.b, 'ok')
	self.assertEqual(obj_loaded.my_property, 'okok')