add init.py

Author: A9isha

src/maxtext/__init__.py

@@ -0,0 +1,45 @@
+# Copyright 2023–2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+MaxText is a high performance, highly scalable, open-source LLM written in pure Python/Jax and targeting Google Cloud
+TPUs and GPUs for training and inference. MaxText achieves high MFUs and scales from single host to very large clusters
+while staying simple and "optimization-free" thanks to the power of Jax and the XLA compiler.
+"""
+
+__author__ = "Google LLC"
+__version__ = "0.2.0"
+__description__ = (
+    "MaxText is a high performance, highly scalable, open-source LLM written in pure Python/Jax and "
+    "targeting Google Cloud TPUs and GPUs for training and **inference."
+)
+
+from collections.abc import Sequence
+
+import os
+# In order to have any effect on the C++ logging this has to be set before we import anything from jax.
+# When jax is imported, its `__init__.py` calls `cloud_tpu_init()`, which also initializes the C++ logger.
+os.environ.setdefault("TF_CPP_MIN_LOG_LEVEL", "0")
+del os
+
+from jax.sharding import Mesh
+
+from maxtext.configs import pyconfig
+from maxtext.models import models
+from maxtext.trainers.post_train.dpo import dpo_utils
+from maxtext.utils import maxtext_utils
+from maxtext.utils.model_creation_utils import *
+
+Transformer = models.Transformer
+transformer_as_linen = models.transformer_as_linen