feat(train): use compilation cache

tools/train/train.py

@@ -41,6 +41,7 @@
 from flax.training import train_state
 from flax.training.common_utils import onehot
 from jax.experimental import PartitionSpec, maps
+from jax.experimental.compilation_cache import compilation_cache as cc
 from jax.experimental.pjit import pjit, with_sharding_constraint
 from tqdm import tqdm
 from transformers import HfArgumentParser
@@ -53,6 +54,11 @@
     set_partitions,
 )
 
+cc.initialize_cache(
+    "/home/boris/dalle-mini/jax_cache", max_cache_size_bytes=5 * 2**30
+)
+
+
 logger = logging.getLogger(__name__)