Remove unused import (AI-Hypercomputer#21)

remove unused import

benchmarks/run_offline.py

@@ -14,17 +14,12 @@
 
 from absl import app
 from absl import flags
-from absl import logging
-import sys
 import jax
 import jax.numpy as jnp
-import numpy as np
 
 from jetstream.engine import token_utils
-from absl.testing import absltest
 
 import os
-import sys
 
 from jetstream_pt import engine as je
 import time

jetstream_pt/engine.py

@@ -14,18 +14,14 @@
 
 """Implement Jet Engine API."""
 
-import copy
 from typing import Any, List, Optional, Tuple, Union
 import threading
 import functools
 
 from flax import struct
-from absl import logging
 import jax
 from jax import numpy as jnp
-from jax.experimental import mesh_utils
 import torch
-import jax.sharding as jsharding
 import numpy as np
 
 from jetstream.engine import engine_api, tokenizer_pb2, token_utils
@@ -37,7 +33,6 @@
 from jetstream_pt.environment import JetEngineEnvironment, JetEngineEnvironmentData
 
 from torch.utils import _pytree as pytree
-import torch
 
 
 

jetstream_pt/layers.py

@@ -16,7 +16,7 @@
 """This version contains modification to make it easier to trace and support batch."""
 
 import math
-from typing import Any, List, Optional, Tuple
+from typing import Optional, Tuple
 
 import torch
 from torch import nn

jetstream_pt/third_party/llama2/model_exportable.py

@@ -1,20 +1,14 @@
 # pylint: disable-all
 """This version contains modification to make it easier to trace and support batch."""
 
-import math
-from typing import Any, List, Optional, Tuple
+from typing import Any, List, Optional
 
 import torch
 from torch import nn
 import torch.nn.functional as F
 
 from . import model_args 
-import jax.sharding as jsharding
-from jax.experimental import mesh_utils
 import jax
-import jax.numpy as jnp
-import torch_xla2
-import torch_xla2.extra
 
 from jetstream_pt.layers import Attention, RMSNorm, Int8Embedding, WeightOnlyInt8Linear
 

jetstream_pt/third_party/llama2/tokenizer.py

@@ -2,7 +2,6 @@
 # This software may be used and distributed according to the terms of the Llama 2 Community License Agreement.
 
 import os
-from logging import getLogger
 from typing import List
 
 from sentencepiece import SentencePieceProcessor

run_interactive.py

@@ -17,22 +17,16 @@
 from absl import logging
 import random
 from typing import List
-import sys
 import jax
-import jax.numpy as jnp
-import numpy as np
 
 from jetstream.engine import token_utils
-from absl.testing import absltest
-from colorama import Fore, Back, Style
+from colorama import Fore, Style
 
 
 import os
-import sys
 
 from jetstream_pt import engine as je
 import time
-import logging
 
 
 logging.getLogger().setLevel(logging.ERROR)

run_server.py

@@ -21,7 +21,6 @@
 
 from jetstream.core import server_lib
 import jetstream_pt
-from jetstream_pt import config
 from jetstream.core.config_lib import ServerConfig