Describe the bug
Please avoid have default values as lambda function. They cause issues with serialization and pickling.
Also theses instances cannot be used with asyncio.
Note that the tensorflow classes do not have this issues and they can safely used with asyncio.
Instead of using lambda as defaults have them as global functions.
Code to reproduce the issue
Please check the code here
Line 238 and Line 312
The scale_fn is lambda function. Can you please define them as global function. Defining it like this creates local functions and thus breaks serialization.