Added support for python copy (#335)

* Added support for python copy

* precommit changes

* removed `_compiled_call_impl` line

* added tests and suggested changes

* ACK changes

python/mlx/nn/layers/base.py

@@ -530,3 +530,9 @@ class Module(dict):
         See :func:`train`.
         """
         self.train(False)
+
+    def __getstate__(self):
+        return self.__dict__.copy()
+
+    def __setstate__(self, state):
+        self.__dict__.update(state)