Make MaxText as Python Modules

[JoeZijunZhou]

Can we refactor the imports to make MaxText as Python Modules? It's pretty hard for developers to use or develop on top of it.

• Blocking inference development with JetStream.
• Hard to integrate with other python libraries or modules.

[shralex]

This was was addressed. It is now possible.