Viewing a single comment thread. View all comments

patrickkidger t1_ir8dz9i wrote

Reply to comment by Small_Stand_8716 in [R] Introduction to Diffusion Models in JAX by jakiwjakiw

Obligatory advertisement for Equinox as a third alternative.

(And IMO much easier to use, although I am biased.)

2

nnexx_ t1_irbf3qt wrote

Yes equinox is amazing. Also obligatory mention of diffrax that implements ODE SDE and CDE in equinox. Diffrax’s documentation even has a working continuous time diffusion model code example !

3

Small_Stand_8716 t1_irajgt0 wrote

Haha, yes, I've stumbled upon Equinox many times before during my JAX research. It is far more appealing to me, syntax wise, than Haiku or Flax and has the ease of use of PyTorch. Its only downside, in my opinion, is that it doesn't have as robust an ecosystem, e.g., pre-trained vision model. Excellent package nevertheless!

2

patrickkidger t1_irake85 wrote

Have you seen eqxvision? It's still relatively nascent but already has a fair amount.

3

Small_Stand_8716 t1_ire7h35 wrote

Thank you, eqxvision fits the bill for me actually. For the most part, I need pre-trained backbones (usually just RegNet, EfficientNetV2, or ViT) and code the rest myself (YOLO, DETR, and so forth). Thanks again!

1