Refactoring to make the polymorphic code easier to read.

Changes:
1. Remove "self._created_variables = []" as this shouldn't matter.
2. Remove self._concrete_stateful_fn -> this is used only for the first time we call __call__.
3. Make _initialize not return canonicalized arguments, instead create a function for it.
PiperOrigin-RevId: 225794015