Increment version

Author: dirmeier

.pre-commit-config.yaml

@@ -10,8 +10,6 @@ repos:
     - id: check-yaml
     - id: debug-statements
     - id: end-of-file-fixer
-    - id: no-commit-to-branch
-      args: [--branch, main]
     - id: requirements-txt-fixer
     - id: trailing-whitespace
 
@@ -20,7 +18,7 @@ repos:
   hooks:
   - id: mypy
     args: ["--ignore-missing-imports"]
-    files: "(fll)"
+    files: "(blaxbird)"
 
 - repo: https://github.com/astral-sh/ruff-pre-commit
   rev: v0.3.0

examples/cifar10_flow_matching/main.py

@@ -1,6 +1,7 @@
 import argparse
 import os
 
+import dataloader
 import jax
 import matplotlib.pyplot as plt
 import numpy as np
@@ -12,7 +13,6 @@
 from jax.experimental import mesh_utils
 
 import blaxbird
-import dataloader
 from blaxbird import get_default_checkpointer, train_fn
 from blaxbird.experimental import rfm
 
@@ -38,19 +38,19 @@ def visualize_hook(sample_fn, val_iter, hook_every_n_steps, log_to_wandb):
 
   def convert_batch_to_image_grid(image_batch):
     reshaped = (
-        image_batch.reshape(n_row, n_col, *img_size)
-        .transpose([0, 2, 1, 3, 4])
-        .reshape(n_row * img_size[0], n_col * img_size[1], img_size[2])
+      image_batch.reshape(n_row, n_col, *img_size)
+      .transpose([0, 2, 1, 3, 4])
+      .reshape(n_row * img_size[0], n_col * img_size[1], img_size[2])
     )
     return (reshaped + 1.0) / 2.0
 
   def plot(images):
     fig = plt.figure(figsize=(16, 6))
     ax = fig.add_subplot(1, 1, 1)
     ax.imshow(
-        images,
-        interpolation="nearest",
-        cmap="gray",
+      images,
+      interpolation="nearest",
+      cmap="gray",
     )
     plt.axis("off")
     plt.tight_layout()
@@ -61,26 +61,31 @@ def fn(step, *, model, **kwargs):
       return
     all_samples = []
     for i, batch in enumerate(val_iter):
-        samples = sample_fn(model, jr.fold_in(jr.key(step), i), sample_shape=batch["inputs"].shape)
-        all_samples.append(samples)
-        if len(all_samples) * all_samples[0].shape[0] >= n_row * n_col:
-          break
-    all_samples = np.concatenate(all_samples, axis=0)[:(n_row * n_col)]
+      samples = sample_fn(
+        model, jr.fold_in(jr.key(step), i), sample_shape=batch["inputs"].shape
+      )
+      all_samples.append(samples)
+      if len(all_samples) * all_samples[0].shape[0] >= n_row * n_col:
+        break
+    all_samples = np.concatenate(all_samples, axis=0)[: (n_row * n_col)]
     all_samples = convert_batch_to_image_grid(all_samples)
     fig = plot(all_samples)
     if jax.process_index() == 0 and log_to_wandb:
-       wandb.log({"images": wandb.Image(fig)}, step=step)
+      wandb.log({"images": wandb.Image(fig)}, step=step)
 
   return fn
 
 
-def get_hooks(sample_fn, val_itr, hook_every_n_steps, log_to_wandb ):
+def get_hooks(sample_fn, val_itr, hook_every_n_steps, log_to_wandb):
   return [visualize_hook(sample_fn, val_itr, hook_every_n_steps, log_to_wandb)]
 
 
 def get_train_and_val_itrs(rng_key, outfolder):
   return dataloader.data_loaders(
-    rng_key, outfolder, split=["train[:90%]", "train[90%:]"], shuffle=[True, False],
+    rng_key,
+    outfolder,
+    split=["train[:90%]", "train[90%:]"],
+    shuffle=[True, False],
   )
 
 
@@ -92,14 +97,19 @@ def run(n_steps, eval_every_n_steps, n_eval_batches, dit_type, log_to_wandb):
     jr.key(0), os.path.join(outfolder, "data")
   )
 
-  model = getattr(blaxbird.experimental, dit_type)(image_size=(32, 32, 3), rngs=nnx.rnglib.Rngs(jr.key(1)))
+  model = getattr(blaxbird.experimental, dit_type)(
+    image_size=(32, 32, 3), rngs=nnx.rnglib.Rngs(jr.key(1))
+  )
   train_step, val_step, sample_fn = rfm()
   optimizer = get_optimizer(model)
 
   save_fn, _, restore_last_fn = get_default_checkpointer(
-    os.path.join(outfolder, "checkpoints"), save_every_n_steps=eval_every_n_steps
+    os.path.join(outfolder, "checkpoints"),
+    save_every_n_steps=eval_every_n_steps,
   )
-  hooks = get_hooks(sample_fn, val_itr, eval_every_n_steps, log_to_wandb) + [save_fn]
+  hooks = get_hooks(sample_fn, val_itr, eval_every_n_steps, log_to_wandb) + [
+    save_fn
+  ]
 
   model_sharding, data_sharding = get_sharding()
   model, optimizer = restore_last_fn(model, optimizer)
@@ -121,7 +131,15 @@ def run(n_steps, eval_every_n_steps, n_eval_batches, dit_type, log_to_wandb):
   parser.add_argument("--n-steps", type=int, default=1_000)
   parser.add_argument("--eval-every-n-steps", type=int, default=50)
   parser.add_argument("--n-eval-batches", type=int, default=10)
-  parser.add_argument("--dit", type=str, choices=["SmallDiT", "BaseDiT"], default="SmallDiT")
+  parser.add_argument(
+    "--dit", type=str, choices=["SmallDiT", "BaseDiT"], default="SmallDiT"
+  )
   parser.add_argument("--log-to-wandb", action="store_true")
   args = parser.parse_args()
-  run(args.n_steps, args.eval_every_n_steps, args.n_eval_batches, args.dit, args.log_to_wandb)
+  run(
+    args.n_steps,
+    args.eval_every_n_steps,
+    args.n_eval_batches,
+    args.dit,
+    args.log_to_wandb,
+  )

examples/mnist_classification/main.py

@@ -8,11 +8,10 @@
 from flax import nnx
 from jax import random as jr
 from jax.experimental import mesh_utils
+from model import CNN, train_step, val_step
 
 from blaxbird import get_default_checkpointer, train_fn
 
-from model import CNN, train_step, val_step
-
 
 def get_optimizer(model, lr=1e-4):
   tx = optax.adamw(lr)
@@ -83,7 +82,8 @@ def run(n_steps, eval_every_n_steps, n_eval_batches):
   optimizer = get_optimizer(model)
 
   save_fn, _, restore_last_fn = get_default_checkpointer(
-    os.path.join(outfolder, "checkpoints"), save_every_n_steps=eval_every_n_steps
+    os.path.join(outfolder, "checkpoints"),
+    save_every_n_steps=eval_every_n_steps,
   )
   hooks = get_hooks(val_itr, eval_every_n_steps) + [save_fn]
 

pyproject.toml

@@ -68,14 +68,15 @@ addopts = "-v --doctest-modules --cov=./blaxbird --cov-report=xml"
 [tool.ruff]
 indent-width = 2
 line-length = 80
-exclude = ["*_test.py", "docs/**", "examples/**"]
 
 [tool.ruff.lint]
 select = ["D", "E", "F", "W", "I001"]
 extend-select = [
   "UP", "I", "PL", "S"
 ]
 ignore = ["S101", "ANN101", "PLR2044", "PLR0913"]
+exclude = ["*_test.py", "docs/**", "examples/**"]
+
 
 [tool.ruff.lint.pydocstyle]
 convention= 'google'