Get all tests to pass locally with no special configuration

MaxText/convert_gemma2_chkpt.py

@@ -33,7 +33,7 @@
 import orbax
 
 import checkpointing
-from train import save_checkpoint
+from MaxText.train import save_checkpoint
 
 Params = dict[str, Any]
 

MaxText/convert_gemma_chkpt.py

@@ -33,7 +33,7 @@
 import orbax
 
 import checkpointing
-from train import save_checkpoint
+from MaxText.train import save_checkpoint
 
 Params = dict[str, Any]
 

MaxText/convert_gpt3_ckpt_from_paxml.py

@@ -32,7 +32,7 @@
   --run-name=$RUN_NAME \
   --base-output-directory=$BASE_OUTPUT_DIR
 """
-import max_utils
+from MaxText import max_utils
 import optimizers
 import pyconfig
 import os
@@ -50,7 +50,7 @@
 import gc
 import max_logging
 from psutil import Process
-from train import save_checkpoint
+from MaxText.train import save_checkpoint
 import argparse
 
 

MaxText/decode.py

@@ -16,8 +16,7 @@
 
 import jax
 
-import max_utils
-import maxengine
+from MaxText import max_utils, maxengine
 
 import os
 import pyconfig

MaxText/generate_param_only_checkpoint.py

@@ -25,9 +25,7 @@
 
 import checkpointing
 import jax
-import max_logging
-import max_utils
-import optimizers
+from MaxText import max_logging, max_utils, optimizers
 import pyconfig
 
 from absl import app
@@ -36,7 +34,7 @@
 from jax import random
 from typing import Sequence
 from layers import models, quantizations
-from train import save_checkpoint
+from MaxText.train import save_checkpoint
 
 Transformer = models.Transformer
 

MaxText/inference_microbenchmark.py

@@ -24,7 +24,7 @@
 
 from jetstream.engine import token_utils
 
-import max_utils
+from MaxText import max_utils
 import maxengine
 import maxtext_utils
 import profiler

MaxText/input_pipeline/_tfds_data_processing.py

@@ -24,10 +24,10 @@
 import tensorflow_datasets as tfds
 import jax
 
-import multihost_dataloading
-import tokenizer
-import sequence_packing
-from input_pipeline import _input_pipeline_utils
+from MaxText import multihost_dataloading
+from MaxText import tokenizer
+from MaxText import sequence_packing
+from MaxText.input_pipeline import _input_pipeline_utils
 
 AUTOTUNE = tf.data.experimental.AUTOTUNE
 

MaxText/kernels/ragged_attention.py

@@ -24,7 +24,7 @@
 from jax.experimental.pallas import tpu as pltpu
 import jax.numpy as jnp
 import numpy as np
-import common_types
+from MaxText import common_types
 
 from jax.experimental import shard_map
 

MaxText/layers/attentions.py

@@ -29,13 +29,13 @@
 from jax.experimental.pallas.ops.tpu.splash_attention import splash_attention_kernel
 from jax.experimental.pallas.ops.tpu.splash_attention import splash_attention_mask
 import jax.numpy as jnp
-import common_types
-from kernels.ragged_attention import ragged_gqa
-from kernels.ragged_attention import ragged_mha
-from layers import embeddings
-from layers import initializers
-from layers import linears
-from layers import quantizations
+from MaxText import common_types
+from MaxText.kernels.ragged_attention import ragged_gqa
+from MaxText.kernels.ragged_attention import ragged_mha
+from MaxText.layers import embeddings
+from MaxText.layers import initializers
+from MaxText.layers import linears
+from MaxText.layers import quantizations
 
 
 # pylint: disable=line-too-long, g-doc-args, g-doc-return-or-yield, bad-continuation, g-inconsistent-quotes

MaxText/layers/embeddings.py

@@ -21,7 +21,7 @@
 import jax
 from jax import lax
 import jax.numpy as jnp
-from layers import initializers
+from MaxText.layers import initializers
 
 Config = Any
 Array = jnp.ndarray

MaxText/layers/initializers.py

@@ -18,7 +18,7 @@
 
 from flax import linen as nn
 import jax
-import common_types
+from MaxText import common_types
 
 Array = common_types.Array
 DType = common_types.DType

MaxText/layers/linears.py

@@ -23,18 +23,18 @@
 import jax
 from jax import lax
 import jax.numpy as jnp
-import common_types
-from layers import initializers
-from layers import normalizations
-from layers import quantizations
+from MaxText import common_types
+from MaxText.layers import initializers
+from MaxText.layers import normalizations
+from MaxText.layers import quantizations
 import numpy as np
 from jax.ad_checkpoint import checkpoint_name
 from jax.experimental import shard_map
 import math
-import max_logging
-import max_utils
+from MaxText import max_logging
+from MaxText import max_utils
 from aqt.jax.v2 import aqt_tensor
-from kernels import megablox as mblx
+from MaxText.kernels import megablox as mblx
 
 
 Array = common_types.Array

MaxText/layers/llama2.py

@@ -23,14 +23,14 @@
 import jax.numpy as jnp
 from jax.ad_checkpoint import checkpoint_name
 # from jax.experimental.pallas.ops.tpu import flash_attention
-from layers import attentions
-from layers import embeddings
-from layers import linears
-from layers import normalizations
-from layers import models
-from layers import quantizations
-
-import common_types
+from MaxText.layers import attentions
+from MaxText.layers import embeddings
+from MaxText.layers import linears
+from MaxText.layers import normalizations
+from MaxText.layers import models
+from MaxText.layers import quantizations
+
+from MaxText import common_types
 from typing import Optional
 
 Array = common_types.Array

MaxText/layers/models.py

@@ -24,12 +24,12 @@
 import jax
 import jax.numpy as jnp
 from jax.ad_checkpoint import checkpoint_name
-import common_types
-from layers import attentions
-from layers import embeddings
-from layers import linears
-from layers import normalizations, quantizations
-from layers import pipeline
+from MaxText import common_types
+from MaxText.layers import attentions
+from MaxText.layers import embeddings
+from MaxText.layers import linears
+from MaxText.layers import normalizations, quantizations
+from MaxText.layers import pipeline
 
 Array = common_types.Array
 Config = common_types.Config

MaxText/layers/normalizations.py

@@ -19,7 +19,7 @@
 from flax import linen as nn
 from jax import lax
 import jax.numpy as jnp
-from layers import initializers
+from MaxText.layers import initializers
 
 Initializer = initializers.Initializer
 

MaxText/layers/pipeline.py

@@ -20,7 +20,7 @@
 from jax import numpy as jnp
 from flax.core import meta
 from flax import linen as nn
-import common_types
+from MaxText import common_types
 import functools
 from typing import Any
 

MaxText/layers/quantizations.py

@@ -23,7 +23,7 @@
 from aqt.jax.v2.flax import aqt_flax
 from aqt.jax.v2 import tiled_dot_general
 from aqt.jax.v2 import calibration
-import common_types
+from MaxText import common_types
 from dataclasses import dataclass
 from flax.linen import fp8_ops
 from flax.linen import initializers as flax_initializers

MaxText/llama_ckpt_conversion_inference_only.py

@@ -37,7 +37,7 @@
 import jax
 from flax.training import train_state
 import max_logging
-from train import save_checkpoint
+from MaxText.train import save_checkpoint
 import torch
 import sys
 import os

MaxText/llama_mistral_mixtral_orbax_to_hf.py

@@ -38,14 +38,14 @@
 from absl import app
 import numpy as np
 import pyconfig
-import max_utils
+from MaxText import max_utils
 from jax.sharding import Mesh
 import max_logging
 import checkpointing
 from generate_param_only_checkpoint import _read_train_checkpoint
 import llama_or_mistral_ckpt
 from transformers import LlamaForCausalLM, MistralForCausalLM, AutoModelForCausalLM, AutoConfig
-from max_utils import unpermute_from_match_maxtext_rope
+from MaxText.max_utils import unpermute_from_match_maxtext_rope
 
 
 def reverse_scale(arr, scale):

MaxText/llama_or_mistral_ckpt.py

@@ -49,10 +49,10 @@
 from tqdm import tqdm
 
 import max_logging
-from train import save_checkpoint
+from MaxText.train import save_checkpoint
 import checkpointing
 from safetensors import safe_open
-import max_utils
+from MaxText import max_utils
 
 MODEL_PARAMS_DICT = {
     "llama2-70b": {

MaxText/load_and_quantize_checkpoint.py

@@ -16,8 +16,7 @@
 
 import jax
 
-import max_utils
-import maxengine
+from MaxText import max_utils, maxengine
 
 import os
 import pyconfig

MaxText/max_utils.py

@@ -20,8 +20,7 @@
 import jax
 import jax.numpy as jnp
 from jax.experimental import mesh_utils
-import checkpointing
-import common_types
+from MaxText import checkpointing, common_types
 import functools
 import time
 import optax

MaxText/maxengine.py

@@ -22,22 +22,21 @@
 from flax import linen as nn
 from flax.linen import partitioning as nn_partitioning
 
-from layers import models, quantizations
+from MaxText.layers import models, quantizations
 
 import jax
 import jax.numpy as jnp
 from jax.sharding import PartitionSpec as P
 from jax.experimental import layout as jax_layout
 
-import common_types
+from MaxText import common_types
 from jetstream.core import config_lib
 from jetstream.engine import engine_api
 from jetstream.engine import tokenizer_pb2
 from jetstream.engine import tokenizer_api
 from jetstream.engine import token_utils
 
-import max_utils
-import inference_utils
+from MaxText import max_utils, inference_utils
 import pyconfig
 
 import warnings

MaxText/maxtext_utils.py

@@ -19,7 +19,7 @@
 
 import jax
 import optax
-import max_utils
+from MaxText import max_utils
 from jax.sharding import PartitionSpec as P
 from jax.experimental.serialize_executable import deserialize_and_load
 

MaxText/metric_logger.py

@@ -22,8 +22,7 @@
 import os
 import numpy as np
 
-import max_logging
-import max_utils
+from MaxText import max_utils, max_logging
 
 
 def _prepare_metrics_for_json(metrics, step, run_name):

MaxText/pyconfig.py

@@ -26,8 +26,7 @@
 from jax.experimental.compilation_cache import compilation_cache
 from layers.attentions import AttentionType
 import accelerator_to_spec_map
-import max_logging
-import max_utils
+from MaxText import max_utils, max_logging
 import omegaconf
 
 OmegaConf = omegaconf.OmegaConf

MaxText/standalone_checkpointer.py

@@ -30,11 +30,10 @@
 from jax import numpy as jnp
 import numpy as np
 
-import checkpointing
-import max_utils
-import max_logging
+from MaxText import checkpointing
+from MaxText import max_utils, max_logging
 import pyconfig
-from train import setup_mesh_and_model, get_first_step, validate_train_config, save_checkpoint
+from MaxText.train import setup_mesh_and_model, get_first_step, validate_train_config, save_checkpoint
 
 from layers import models
 

MaxText/standalone_dataloader.py

@@ -27,7 +27,7 @@
 import numpy as np
 
 import pyconfig
-from train import validate_train_config, get_first_step, load_next_batch, setup_train_loop
+from MaxText.train import validate_train_config, get_first_step, load_next_batch, setup_train_loop
 
 
 def data_load_loop(config, state=None):

MaxText/tests/attention_test.py

@@ -20,17 +20,17 @@
 import unittest
 from absl.testing import parameterized
 
-import common_types
-
+from MaxText import common_types, max_utils
 from flax.core import freeze
 import jax
 import jax.numpy as jnp
+
 import max_utils
 import pytest
 
 import pyconfig
 
-from layers import attentions
+from MaxText.layers import attentions
 
 Mesh = jax.sharding.Mesh
 Attention = attentions.Attention