Repository: google/flax

Issue:
Exceptions are not pickle-able
### Problem you have encountered:

Flax exceptions can not be pickled. This makes it impossible to trace where/why errors are occurring when calling flax via ray.


### System information
```
Name: flax
Version: 0.8.3
Summary: Flax: A neural network library for JAX designed for flexibility
Home-page:
Author:
Author-email: Flax team <flax-dev@google.com>
License:
Location: /home/sander_cohere_com/miniconda3/envs/ct/lib/python3.10/site-packages
Requires: jax, msgpack, numpy, optax, orbax-checkpoint, PyYAML, rich, tensorstore, typing-extensions
Required-by: dm-haiku, fax
---
Name: jax
Version: 0.4.28
Summary: Differentiate, compile, and transform Numpy code.
Home-page: https://github.com/google/jax
Author: JAX team
Author-email: jax-dev@google.com
License: Apache-2.0
Location: /home/sander_cohere_com/miniconda3/envs/ct/lib/python3.10/site-packages
Requires: ml-dtypes, numpy, opt-einsum, scipy
Required-by: chex, fax, flax, optax, orbax-checkpoint, rax
---
Name: jaxlib
Version: 0.4.28
Summary: XLA library for JAX
Home-page: https://github.com/google/jax
Author: JAX team
Author-email: jax-dev@google.com
License: Apache-2.0
Location: /home/sander_cohere_com/miniconda3/envs/ct/lib/python3.10/site-packages
Requires: ml-dtypes, numpy, scipy
Required-by: chex, fax, optax, orbax-checkpoint, rax
```

### What you expected to happen:

Flax Exceptions to support pickle, cf [this guide](https://stackoverflow.com/questions/16244923/how-to-make-a-custom-exception-class-with-multiple-init-args-pickleable)

### Logs, error messages, etc:

In production code:
```
site-packages/ray/exceptions.py", line 49, in from_ray_exception
    return pickle.loads(ray_exception.serialized_exception)
TypeError: ScopeParamShapeError.__init__() missing 3 required positional arguments: 'scope_path', 'value_shape', and 'init_shape'
```


### Steps to reproduce:

```
import flax.linen as nn
from jax import random
from flax.linen.initializers import lecun_normal
from jax import lax
import pickle

class NoBiasDense(nn.Module):
  features: int = 8

  @nn.compact
  def __call__(self, x):
    kernel = self.param('kernel',
                        lecun_normal(),
                        (x.shape[-1], self.features))  # <--- Exception from flax docs example
    y = lax.dot_general(x, kernel,
                        (((x.ndim - 1,), (0,)), ((), ())))
    return y

variables = NoBiasDense().init(random.key(0), jnp.ones((5, 5, 1)))
try:
    _ = NoBiasDense().apply(variables, jnp.ones((5, 5)))
except Exception as e:
    str = pickle.dumps(e)
    obj = pickle.loads(str) # <--- pickle exception
```


Discussion:

Patch:
diff --git a/flax/errors.py b/flax/errors.py
index 1308af3a..8863ad06 100644
--- a/flax/errors.py
+++ b/flax/errors.py
@@ -56,8 +56,12 @@ class FlaxError(Exception):
     )
     module_name = self.__class__.__module__
     class_name = self.__class__.__name__
-    error_msg = f'{message} ({error_page}#{module_name}.{class_name})'
-    super().__init__(error_msg)
+    if error_page not in message: # do not add a FlaxError link on unpickling
+      message = f'{message} ({error_page}#{module_name}.{class_name})'
+    super().__init__(message)
+
+  def __reduce__(self):
+   return (FlaxError, (str(self),))
 
 
 #################################################


Test patch:
diff --git a/tests/pickle_test.py b/tests/pickle_test.py
new file mode 100644
index 00000000..e65e7524
--- /dev/null
+++ b/tests/pickle_test.py
@@ -0,0 +1,34 @@
+# Copyright 2024 The Flax Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for flax.errors."""
+
+from absl.testing import absltest
+from flax.errors import FlaxError, ScopeVariableNotFoundError
+import pickle
+
+class ErrorrsTest(absltest.TestCase):
+  def test_exception_can_be_pickled(self):
+    # tests the new __reduce__ method fixes bug reported in issue #4000
+    ex = ScopeVariableNotFoundError('varname', 'collection', 'scope')
+    pickled_ex = pickle.dumps(ex)
+    unpicked_ex = pickle.loads(pickled_ex)
+    self.assertIsInstance(unpicked_ex, FlaxError)
+    self.assertIn('varname', str(unpicked_ex))
+    self.assertIn('#flax.errors.ScopeVariableNotFoundError', str(unpicked_ex))
+    self.assertNotIn('#flax.errors.FlaxError', str(unpicked_ex))
+
+
+if __name__ == '__main__':
+  absltest.main()