tensorflow - 如何将 Tensorflow1 代码转换为 JAX 代码

Question

我正在尝试将张量流代码转换为 JAX 代码。我的困难是 Stackoverflow 中几乎没有关于 JAX 的任何材料。以下是我要转换的代码，任何帮助将不胜感激。

tf.reset_default_graph()
X = tf.placeholder(tf.float32, [n_dim, None])
Y = tf.placeholder(tf.float32, [1, None])
learning_rate = tf.placeholder(tf.float32, shape=())
W = tf.Variable(tf.ones([n_dim,1]))
b = tf.Variable(tf.zeros(1))
init = tf.global_variables_initializer()
y_ = tf.matmul(tf.transpose(W),X)+b
cost = tf.reduce_mean(tf.square(y_-Y))
training_step = tf.train.GradientDescentOptimizer(learning_rate).minimize(cost)