Cleanup - #6

Merged

dirmeier merged 34 commits into

mainfrom

fleanup

Jul 14, 2026

Commits on Jul 9, 2026

chore: gitignore .worktrees/
dirmeier
committed
e028f2a

Commits on Jul 14, 2026

refactor: move model zoo (DiT, MLP, edm, rfm, samplers) from blaxbird package to examples/_common
dirmeier
committed
2b09b11
refactor(samplers): add SAMPLERS registry and get_sampler_fn
dirmeier
committed
3d01169
refactor(parameterizations): add EDMParameterization.denoise
dirmeier
committed
26d2cf1
feat(types): add ObjectiveFns contract, wire edm/rfm to return it
dirmeier
committed
a64c8d1
refactor(trainer): derive model from optimizer.model, drop redundant param
dirmeier
committed
c3930b2
refactor(checkpointer): restore_best_fn/restore_last_fn take/return optimizer only
dirmeier
committed
3dd82a8
feat(checkpointer): make checkpoint criterion_key and best_mode configurable
dirmeier
committed
c0b71eb
fix(dit): actually use context for class conditioning instead of silently dropping it
dirmeier
committed
2322ebb
feat(examples): class-conditional CIFAR-10 generation, exercising the DiT context fix
dirmeier
committed
745c521
feat(unet): add ResBlock, AttentionBlock, Downsample, Upsample building blocks
dirmeier
committed
bf6f553
feat(unet): assemble full UNet with down/mid/up path and Small/Base/Large presets
dirmeier
committed
6f5bc66
feat(_common): export UNet and presets from _common
dirmeier
committed
54fd6ab
test(unet): verify gradients flow through the full network
dirmeier
committed
f02a77b
feat(hooks): add get_ema_hook, first entry in the new blaxbird.hooks module
dirmeier
committed
1b3f418
feat: export get_ema_hook from blaxbird
dirmeier
committed
95e3556
test(hooks): verify get_ema_hook integrates with train_fn's hooks=
dirmeier
committed
9af82e7
feat(examples): add LR warmup/cosine-decay schedule and gradient clipping to cifar10 optimizer
dirmeier
committed
efd2170
feat(trainer): derive sharding from nnx.with_partitioning via mesh, replacing single blanket shardings arg
dirmeier
committed
dc5559e
feat(examples): add FSDP+TP 2D mesh sharding demo
dirmeier
committed
b1cb0f9
docs: update README for mesh-based sharding and optimizer-only train()/restore APIs
dirmeier
committed
fdec12b
feat(llm_reference): add shared RoPE, masking, RMSNorm, TP-sharded GQAAttention/GeGLU
dirmeier
committed
76ba103
feat(llm_reference): assemble GemmaDense with interleaved local/global attention
dirmeier
committed
58f6662
feat(llm_reference): add MLAAttention with decoupled RoPE
dirmeier
committed
75140a4
feat(llm_reference): assemble DeepSeekMLA, verify 2D fsdp+tp sharding
dirmeier
committed
e6f7bd6
feat(llm_reference): add SparseMoEFFN with capacity-based dispatch/combine routing
dirmeier
committed
26ef1f3
feat(llm_reference): assemble MixtralSMoE, verify 3D fsdp+tp+expert sharding matches unsharded reference
dirmeier
committed
5c1356f
feat(llm_reference): add causal_lm training objective, shared across all three model families
dirmeier
committed
22319d2
feat(llm_reference): add shared autoregressive generation loop
dirmeier
committed
a084e99
docs(llm_reference): add multi-mesh smoke-test script and README
dirmeier
committed
5451aa0
feat(llm): add Gemma4/DeepSeek4/Qwen3Next reference architectures
dirmeier
committed
f6b6bbd
chore: reorg examples, drop unused code, make wandb optional
dirmeier
committed
f9e9c1d
feat: add lms
dirmeier
committed
50ca3a0
chore: bump version to 0.2.0
dirmeier
committed
792a6e3

Navigation Menu

Cleanup - #6

Cleanup#6
dirmeier merged 34 commits into
mainfrom
fleanup

Commits on Jul 9, 2026

chore: gitignore .worktrees/

Commits on Jul 14, 2026

refactor: move model zoo (DiT, MLP, edm, rfm, samplers) from blaxbird package to examples/_common

refactor(samplers): add SAMPLERS registry and get_sampler_fn

refactor(parameterizations): add EDMParameterization.denoise

feat(types): add ObjectiveFns contract, wire edm/rfm to return it

refactor(trainer): derive model from optimizer.model, drop redundant param

refactor(checkpointer): restore_best_fn/restore_last_fn take/return optimizer only

feat(checkpointer): make checkpoint criterion_key and best_mode configurable

fix(dit): actually use context for class conditioning instead of silently dropping it

feat(examples): class-conditional CIFAR-10 generation, exercising the DiT context fix

feat(unet): add ResBlock, AttentionBlock, Downsample, Upsample building blocks

feat(unet): assemble full UNet with down/mid/up path and Small/Base/Large presets

feat(_common): export UNet and presets from _common

test(unet): verify gradients flow through the full network

feat(hooks): add get_ema_hook, first entry in the new blaxbird.hooks module

feat: export get_ema_hook from blaxbird

test(hooks): verify get_ema_hook integrates with train_fn's hooks=

feat(examples): add LR warmup/cosine-decay schedule and gradient clipping to cifar10 optimizer

feat(trainer): derive sharding from nnx.with_partitioning via mesh, replacing single blanket shardings arg

feat(examples): add FSDP+TP 2D mesh sharding demo

docs: update README for mesh-based sharding and optimizer-only train()/restore APIs

feat(llm_reference): add shared RoPE, masking, RMSNorm, TP-sharded GQAAttention/GeGLU

feat(llm_reference): assemble GemmaDense with interleaved local/global attention

feat(llm_reference): add MLAAttention with decoupled RoPE

feat(llm_reference): assemble DeepSeekMLA, verify 2D fsdp+tp sharding

feat(llm_reference): add SparseMoEFFN with capacity-based dispatch/combine routing

feat(llm_reference): assemble MixtralSMoE, verify 3D fsdp+tp+expert sharding matches unsharded reference

feat(llm_reference): add causal_lm training objective, shared across all three model families

feat(llm_reference): add shared autoregressive generation loop

docs(llm_reference): add multi-mesh smoke-test script and README

feat(llm): add Gemma4/DeepSeek4/Qwen3Next reference architectures

chore: reorg examples, drop unused code, make wandb optional

feat: add lms

chore: bump version to 0.2.0