style: isort

tools/train/distributed_shampoo.py

@@ -36,13 +36,13 @@ import itertools
 from typing import Any, List, NamedTuple
 
 import chex
-from flax import struct
 import jax
-from jax import lax
 import jax.experimental.pjit as pjit
 import jax.numpy as jnp
 import numpy as np
 import optax
+from flax import struct
+from jax import lax
 
 
 # pylint:disable=no-value-for-parameter

tools/train/train.py

@@ -34,6 +34,7 @@ import optax
 import transformers
 import wandb
 from datasets import Dataset
+from distributed_shampoo import GraftingType, distributed_shampoo
 from flax import jax_utils, traverse_util
 from flax.jax_utils import unreplicate
 from flax.serialization import from_bytes, to_bytes
@@ -45,8 +46,6 @@ from transformers import AutoTokenizer, HfArgumentParser
 from dalle_mini.data import Dataset
 from dalle_mini.model import DalleBart, DalleBartConfig
 
-from distributed_shampoo import distributed_shampoo, GraftingType
-
 logger = logging.getLogger(__name__)