gpt4 book ai didi

Is there a module to convert a tensorflow NN to Jax?(是否有将TensorFlow NN转换为JAX的模块?)

转载 作者:bug小助手 更新时间:2023-10-24 21:44:52 31 4
gpt4 key购买 nike



There is a libary to convert Jax functions to Tensorflow functions. Is there a similar library to convert TensorFlow functions to Jax functions?

有一个库可以将JAX函数转换为TensorFlow函数。是否有类似的库将TensorFlow函数转换为JAX函数?


更多回答

this suggestion isn't possible at the moment, but one could imagine using onnx as an intermediate format. jax --> onnx --> tf. but jax --> onnx doesn't exist (yet).

这个建议目前是不可能的,但可以想象使用onnx作为中间格式。JAX-->ONX-->TF.但JAX-->onnx(目前)还不存在。

@jakub Yep, I had a look at that but couldnt find anything

@Jakub是的,我看过了,但什么也找不到

you can check github.com/deepmind/tf2jax

你可以查看githeb.com/DeepMind/tf2jax

优秀答案推荐

No, there is no library supported by the JAX team to convert tensorflow into JAX in a manner similar to how jax.experimental.jax2tf converts JAX code to tensorflow, and I have not seen any such library developed by others.

不,JAX团队没有支持将TensorFlow转换为JAX的库,其转换方式类似于jax.exestental.jax2tf将JAX代码转换为TensorFlow的方式,我也没有看到其他人开发过这样的库。



See https://github.com/google/jax/blob/main/jax/experimental/jax2tf/README.md.

请参阅https://github.com/google/jax/blob/main/jax/experimental/jax2tf/README.md.


"jax2tf.call_tf: for using TensorFlow functions in a JAX context, e.g., to call a TensorFlow library or a SavedModel inside a JAX function."

Jax2tf.call_tf:用于在JAX上下文中使用TensorFlow函数,例如,在JAX函数中调用TensorFlow库或SavedModel。


That is what you need. So you can call tf function under jax context.

这就是你所需要的。因此,您可以在JAX上下文中调用tf函数。


for example:

例如:


# Compute cos with TF and sin with JAX
def cos_tf_sin_jax(x):
return jax.numpy.sin(jax2tf.call_tf(cos_tf)(x))


Is https://github.com/google-deepmind/tf2jax what you were looking for? It only works for TF v2 though.

您正在寻找的是https://github.com/google-deepmind/tf2jax吗?它只适用于TF V2。



To my knowledge there is no library similar to the one you mentioned to convert TensorFlow functions to Jax functions. I'm sorry

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


更多回答

Do you think there are any hacky ways to allow a tensorflow NN to be used in Jax? I dont need to modify or train the neural network any more, I only need to use it. For instance, could I do something to create a XLA HLO graph and then load it into Jax, within python? I saw this discussion - github.com/google/jax/issues/45 - But it is closed now

你认为有什么破解的方法可以让TensorFlow NN在JAX中使用吗?我不再需要修改或训练神经网络,我只需要使用它。例如,我能做些什么来创建一个XLA HLO图,然后在Python中将其加载到JAX中吗?我看到了这个讨论--githorb.com/Google/jax/Issues/45--但现在已经关闭了

You might have more luck generating some ideas if you ask at github.com/google/jax/discussions

如果你在githorb.com/Google/jax/讨论上询问,你可能会有更多的运气来产生一些想法

31 4 0
Copyright 2021 - 2024 cfsdn All Rights Reserved 蜀ICP备2022000587号
广告合作:1813099741@qq.com 6ren.com