只知道TF和PyTorch还不够，快来看看怎么从PyTorch转向自动微分神器JAX

import torch
class LSTMCell(torch.nn.Module): 
    def __init__(self, in_dim, out_dim): 
        super(LSTMCell, self).__init__() 
        self.weight_ih = torch.nn.Parameter(torch.rand(4*out_dim, in_dim)) 
        self.weight_hh = torch.nn.Parameter(torch.rand(4*out_dim, out_dim)) 
        self.bias = torch.nn.Parameter(torch.zeros(4*out_dim,))  

    def forward(self, inputs, h, c): 
        ifgo = self.weight_ih @ inputs + self.weight_hh @ h + self.bias 
        i, f, g, o = torch.chunk(ifgo, 4) 
        i = torch.sigmoid(i) 
        f = torch.sigmoid(f) 
        g = torch.tanh(g) 
        o = torch.sigmoid(o) 
        new_c = f * c + i * g 
        new_h = o * torch.tanh(new_c) 
        return (new_h, new_c)

class LSTMLM(torch.nn.Module): 
    def __init__(self, vocab_size, dim=17): 
        super().__init__() 
        self.cell = LSTMCell(dim, dim) 
        self.embeddings = torch.nn.Parameter(torch.rand(vocab_size, dim)) 
        self.c_0 = torch.nn.Parameter(torch.zeros(dim))

    @property 
    def hc_0(self): 
        return (torch.tanh(self.c_0), self.c_0)

    def forward(self, seq, hc): 
         loss = torch.tensor(0.) 
          for idx in seq: 
              loss -= torch.log_softmax(self.embeddings @ hc[0], dim=-1)[idx] 
              hc = self.cell(self.embeddings[idx,:], *hc) 
          return loss, hc  

    def greedy_argmax(self, hc, length=6): 
        with torch.no_grad(): 
            idxs = [] 
            for i in range(length): 
                idx = torch.argmax(self.embeddings @ hc[0]) 
                idxs.append(idx.item()) 
                hc = self.cell(self.embeddings[idx,:], *hc) 
        return idxs

torch.manual_seed(0)
# As training data, we will have indices of words/wordpieces/characters,
# we just assume they are tokenized and integerized (toy example obviously).
import jax.numpy as jnp
vocab_size = 43 # prime trick! :)
training_data = jnp.array([4, 8, 15, 16, 23, 42])

lm = LSTMLM(vocab_size=vocab_size)
print("Sample before:", lm.greedy_argmax(lm.hc_0))

bptt_length = 3 # to illustrate hc.detach-ing

for epoch in range(101): 
    hc = lm.hc_0 
    totalloss = 0. 
    for start in range(0, len(training_data), bptt_length): 
        batch = training_data[start:start+bptt_length] 
        loss, (h, c) = lm(batch, hc) 
        hc = (h.detach(), c.detach()) 
        if epoch % 50 == 0: 
            totalloss += loss.item() 
        loss.backward() 
        for name, param in lm.named_parameters(): 
            if param.grad is not None: 
                param.data -= 0.1 * param.grad 
                del param.grad 
     if totalloss: 
         print("Loss:", totalloss)
         
print("Sample after:", lm.greedy_argmax(lm.hc_0))
Sample before: [42, 34, 34, 34, 34, 34]
Loss: 25.953862190246582
Loss: 3.7642268538475037
Loss: 1.9537211656570435
Sample after: [4, 8, 15, 16, 23, 42]

import random
import time
nr_executions = 0

def pure_fn_1(x): 
    return 2 * x
    
def pure_fn_2(xs): 
    ys = [] 
    for x in xs: 
        # Mutating stateful variables *inside* the function is fine! 
        ys.append(2 * x) 
    return ys

def impure_fn_1(xs): 
    # Mutating arguments has lasting consequences outside the function! :( 
    xs.append(sum(xs)) 
    return xs

def impure_fn_2(x): 
    # Very obviously mutating 
    global state is bad... global 
    nr_executions nr_executions += 1 
    return 2 * x

def impure_fn_3(x): 
    # ...but just accessing it is, too, because now the function depends on the 
    # execution context! 
    return nr_executions * x

def impure_fn_4(x): 
    # Things like IO are classic examples of impurity. 
    # All three of the following lines are violations of purity: 
    print("Hello!") 
    user_input = input() 
    execution_time = time.time() 
    return 2 * x

def impure_fn_5(x): 
    # Which constraint does this violate? Both, actually! You access the current 
    # state of randomness *and* advance the number generator! 
    p = random.random() 
    return p * x
Let's see a pure function that JAX operates on: the example from the intro figure.

# (almost) 1-D linear regression
def f(w, x): 
    return w * x

print(f(13., 42.))
546.0

import jax
import jax.numpy as jnp

# Gradient: with respect to weights! JAX uses the first argument by default.
df_dw = jax.grad(f)

def manual_df_dw(w, x): 
    return x
    
assert df_dw(13., 42.) == manual_df_dw(13., 42.)

print(df_dw(13., 42.))
42.0

class LinearRegressor(): 
    def __init__(self, w, b): 
    self.w = w 
    self.b = b 
    
    def predict(self, x): 
        return self.w * x + self.b 
        
    def rms(self, xs: jnp.ndarray, ys: jnp.ndarray): 
        return jnp.sqrt(jnp.sum(jnp.square(self.w * xs + self.b - ys)))
        
my_regressor = LinearRegressor(13., 0.)

# A kind of loss fuction, used for training
xs = jnp.array([42.0])
ys = jnp.array([500.0])
print(my_regressor.rms(xs, ys))

# Prediction for test data
print(my_regressor.predict(42.))
46.0
546.0

def loss_fn(w, b, xs, ys): 
    my_regressor = LinearRegressor(w, b) 
    return my_regressor.rms(xs=xs, ys=ys)
    
# We use argnums=(0, 1) to tell JAX to give us
# gradients wrt first and second parameter.
grad_fn = jax.grad(loss_fn, argnums=(0, 1))

print(loss_fn(13., 0., xs, ys))
print(grad_fn(13., 0., xs, ys))
46.0
(DeviceArray(42., dtype=float32), DeviceArray(1., dtype=float32))

def loss_fn(params, xs, ys): 
    my_regressor = LinearRegressor(params['w'], params['b']) 
    return my_regressor.rms(xs=xs, ys=ys)

grad_fn = jax.grad(loss_fn)

print(loss_fn({'w': 13., 'b': 0.}, xs, ys))
print(grad_fn({'w': 13., 'b': 0.}, xs, ys))
46.0
{'b': DeviceArray(1., dtype=float32), 'w': DeviceArray(42., dtype=float32)}So this already looks nicer! We could write a training loop like this:

params = {'w': 13., 'b': 0.}

for _ in range(15): 
    print(loss_fn(params, xs, ys)) 
    grads = grad_fn(params, xs, ys) 
    for name in params.keys(): 
        params[name] -= 0.002 * grads[name]
        
# Now, predict:
LinearRegressor(params['w'], params['b']).predict(42.)
46.0
42.47003
38.940002
35.410034
31.880066
28.350098
24.820068
21.2901
17.760132
14.230164
10.700165
7.170166
3.6401978
0.110198975
3.4197998
DeviceArray(500.1102, dtype=float32)

def update_combiner(param, grad, lr=0.002): 
    return param - lr * grad
    
params = jax.tree_multimap(update_combiner, params, grads)
# instead of:
# for name in params.keys():
# params[name] -= 0.1 * grads[name]