对比 PyTorch、TensorFlow、JAX、Theano，我发现都在关注两大问题

import aesarafrom aesara import tensor as at
a = at.dscalar("a") # Define placeholders, which have no values.b = at.dscalar("b")
c = a * b              # c now contains the IR of an expression.TTdc = aesara.grad(c, a) # Convert the IR in c into another one, dc
f_dc = aesara.function([a, b], dc) # Convert the IR into a Python function,assert f_dc(1.5, 2.5) == 2.5       # so we can call it.

import tensorflow.compat.v1 as tf # TensorFlow 1.x APIimport numpy as nptf.disable_eager_execution()
a = tf.placeholder(tf.float32, shape=())b = tf.placeholder(tf.float32, shape=())
c = a * bdc = tf.gradients(c, [a], stop_gradients=[a, b])
with tf.compat.v1.Session() as sess:  # TensorFlow has a runtime to execute the IR,   x = np.single(2)                    # so, no converting it into Python code.   y = np.single(3)       print(sess.run(dc, feed_dict={a:x, b:y}))

import torch
a = torch.tensor(1.0, requires_grad=True) # These are not placeholders, but values.b = torch.tensor(2.0)
c = a * b    # Evaluates c and derives the IR of the backward in c.grad_fn_.c.backward() # Executes c.grad_fn_.print(c.grad)

import tensorflow as tf
a = tf.Variable(1.0) # Like PyTorch, these are values, not placehodlers. b = tf.Variable(2.0)
with tf.GradientTape() as tape:  c = a * bdcda = tape.gradient(c, a)print(dcda)

import jax 
a = 2.0b = 3.0jax.grad(jax.lax.mul)(a, b)  # Compute c = a * b w.r.t. a.  The result is b=3. 
jax.jit(jax.grad(jax.lax.mul))(a,b)
jax.experimental.pjit(jax.grad(jax.lax.mul),                       device_mesh(ntpus))(a,b)

jax.make_jaxpr(jax.lax.mul)(2.0, 3.0)  # Returns the IR representing jax.lax.mul(2,3)jax.make_jaxpr(jax.grad(jax.lax.mul))(2.0, 3.0)  # Returns the IR of grad(mul)(2,3)

import torch, functorcha = torch.tensor([2.0])
b = torch.tensor([3.0])functorch.grad(torch.dot)(a, b)

def f(a, b):  return torch.dot(a, b) # Have to wrap the builtin function dot into f. # 必须将内置函数dot转换成f.  print(functorch.make_fx(f)(a, b).code)print(functorch.make_fx(functorch.grad(f))(a, b).code)

from functorch.compile import aot_functionimport torch.fx as fx
def print_compile_fn(fx_module, args):    print(fx_module)    return fx_moduleaot_fn = aot_function(torch.dot, print_compile_fn)aot_fn(a, b)

import torchfrom torch import autograd
x = torch.tensor(1., requires_grad = True)y = 2*x**3 + 8
first_derivative = autograd.grad(y, x, create_graph=True)print(first_derivative)
second_derivative = autograd.grad(first_derivative, x)print(second_derivative)

import tensorflow as tf
x = tf.Variable(1.0)
with tf.GradientTape() as outer_tape:    with tf.GradientTape() as tape:        y = 2*x**3 + 8        dy_dx = tape.gradient(y, x)        print(dy_dx)    d2y_dx2 = outer_tape.gradient(dy_dx, x)    print(d2y_dx2)

def f(a):  return 2*a**3 + 8
print(jax.grad(f)(1.0))print(jax.grad(jax.grad(f))(1.0))

def f1(): return tf.multiply(a, 17)def f2(): return tf.add(b, 23)r = tf.cond(tf.less(a, b), f1, f2)
with tf.compat.v1.Session() as sess:  # TensorFlow has a runtime to execute the IR,  print(sess.run(r, feed_dict={a:x, b:y}))

def g(x, y):    if tf.reduce_any(x < y):        return tf.multiply(x, 17)    return tf.add(y, 23)    converted_g = tf.autograph.to_graph(g)
import inspectprint(inspect.getsource(converted_g))

jax.lax.cond(a < b, lambda : a*17, lambda: b+23)

def factorial(r, x):  return jax.lax.cond(x <= 1.0, lambda: r, lambda: factorial(r*x, x-1))factorial(1.0, 3.0)

import osimport torchfrom torch import distributed as dist
def main():    use_gpu = torch.cuda.is_available()    local_rank = int(os.environ.get("LOCAL_RANK", "0"))    local_world_size = int(os.environ.get("LOCAL_WORLD_SIZE", "0"))    device = torch.device(f"cuda:{local_rank}" if use_gpu else "cpu")    dist.init_distributed(backend="nccl")
    lst = torch.tensor([local_rank + 100]).to(device)
    # placeholder     rlt_lst = [torch.zeros_like(lst) for _ in range(local_world_size)]    dist.all_gather(rlt_lst, lst, async_op=False)    print("After broadcasting:", rlt_lst)

with tf.device('/job:bar/task:0/device:gpu:2'):    # ops created here have the fully specified device above

strategy = tf.distribute.MirroredStrategy() \           if tf.config.list_physical_devices('GPU') \           else tf.distribute.get_strategy()
with strategy.scope():  model = tf.keras.Sequential([tf.keras.layers.Dense(1, input_shape=(1,))])
model.compile(loss='mse', optimizer='sgd')