是否有将 tensorflow NN 转换为 Jax 的模块？

Posted 2023-02-16

【中文标题】是否有将 tensorflow NN 转换为 Jax 的模块？

【英文标题】：Is there a module to convert a tensorflow NN to Jax?

【发布时间】：2022-01-18 04:37:30

【问题描述】：

有一个libary to convert Jax functions to Tensorflow functions。是否有类似的库可以将 TensorFlow 函数转换为 Jax 函数？

【问题讨论】：

这个建议目前是不可能的，但可以想象使用 onnx 作为一种中间格式。 jax --> onnx --> tf。但jax --> onnx 不存在（目前）。

@jakub 是的，我看了一下，但找不到任何东西

【参考方案1】：

没有，JAX 团队不支持将 tensorflow 转换为 JAX 的方式类似于 jax.experimental.jax2tf 将 JAX 代码转换为 tensorflow，我还没有看到其他人开发过任何此类库。

【讨论】：

您是否认为有任何 hacky 方法可以让 tensorflow NN 在 Jax 中使用？我不再需要修改或训练神经网络，我只需要使用它。例如，我可以做一些事情来创建一个 XLA HLO 图，然后在 python 中将它加载到 Jax 中吗？我看到了这个讨论 - github.com/google/jax/issues/45 - 但它现在关闭了

如果您在github.com/google/jax/discussions提问，您可能会更幸运地产生一些想法

【参考方案2】：

据我所知，没有类似于您提到的将 TensorFlow 函数转换为 Jax 函数的库。对不起