问题描述
来自他们的描述：
问题：
我应该选择哪个基于 jax 的库来实现，比如说 DeepSpeech 模型（由 CNN 层 + LSTM 层 + FC 组成）和 ctc-loss？
UPD。
发现了 explanation 与 dm-haiku 开发者的差异：
Flax 包含更多电池，并带有优化器、混合精度和一些训练循环（我听说这些是分离的，您可以根据需要使用多少）。 Haiku 的目标只是解决 NN 模块和状态管理，它将问题的其他部分留给其他库（例如用于优化的 optax）。
Haiku 旨在将 Sonnet（一个 TF NN 库）移植到 JAX。因此，如果（如 DeepMind）您有大量想要在 JAX 中使用的 Sonnet+TF 代码，并且想要尽可能轻松地迁移该代码（在任一方向），那么 Haiku 是更好的选择。
我认为否则归结为个人喜好。在 Alphabet 中，有 100 名研究人员使用每个库，所以我认为你不会出错。在 DeepMind，我们对 Haiku 进行了标准化，因为它对我们很有意义。我建议查看两个库提供的示例代码，看看哪个符合您对结构化实验的偏好。我想如果您将来改变主意，您会发现将代码从一个库迁移到另一个库并不是很复杂。
原始问题仍然相关。
解决方法
我最近遇到了同样的问题，我偏爱俳句，因为我认为它们的实现（参见 Flax Dense() 与 Haiku Linear() ）更接近原始 JAX 精神（即链接 init 和 { {1}} 函数并跟踪 Pytrees 中的参数），这使我可以更轻松地修改内容。
但如果您不想深入修改内容，最好的选择方法是找到一篇关于 CNNs + LSTMs with Flax/Haiku 的不错的博客文章并坚持下去。我的总体看法是，即使我更喜欢 Haiku（+ Optax + Rlax + Chex + ...）的构建方式，但两个库都非常接近。