Adds --ccache (doesn't really help though).

Author: andsteing

examples/controlnet/train_controlnet_flax.py

@@ -22,6 +22,7 @@
 import time
 
 import jax
+import jax.experimental.compilation_cache.compilation_cache as cc
 import jax.numpy as jnp
 import numpy as np
 import optax
@@ -230,6 +231,12 @@ def parse_args():
         action="store_true",
         help="Whether to dump an initial (before training loop) and a final (at program end) memory profile.",
     )
+    parser.add_argument(
+        "--ccache",
+        type=str,
+        default=None,
+        help="Enables compilation cache.",
+    )
     parser.add_argument(
         "--tokenizer_name",
         type=str,
@@ -971,6 +978,9 @@ def cumul_grad_step(grad_idx, loss_grad_rng):
             }
         )
 
+    if args.ccache:
+        cc.initialize_cache(args.ccache)
+
     global_step = 0
     epochs = tqdm(
         range(args.num_train_epochs),