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

我正在尝试将tensorflow代码转换为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)