Custom layer serialization implementation

[seanmor5]

Currently custom layers will fail when they are imported if the implementation is not an MFA function (&Module.function/arity). We should consider raising in these cases, or support some other way of serializing custom layers. Another option is to trace the custom layer and serialize the resulting Nx expression and then rebuild the function from the expression

[josevalim]

Serializing the Nx expression would definitely work... but any dynamic part of the expression would be lost (for example, a case that handles Axon.None structs. In other words, we could serialize it if we have the input shapes.