Informative video, thanks! Although I don't believe this is as groundbreaking as some paint it to be, it is still opens up numerous new avenues to explore.

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!

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!

Thank you for your reply!

I'm not aware of the performance of detach, but why not set requires_grad to False to freeze some layers? It will tremendously speed up training and memory usage.

Great post, thank you! Unrelated, but is there a reason you used Flax over Haiku? I've been meaning to learn JAX for a while and go beyond PyTorch, but I'm unsure whether, for deep learning, I should start with Flax or Haiku.