import jax
from jax import numpy as jnp
from jax import grad, jit, vmap

# 创建 PRNGKey (PRNG State)
key = jax.random.PRNGKey(0)


## 创建模型参数, 去除输入层，实际上三层Linear，每层都包含一组(w, b)，共三组参数

def random_layer_params(m, n, key, scale=1e-2):
    """
    A helper function to randomly initialize weights and biases
    for a dense neural network layer
    """
    w_key, b_key = jax.random.split(key)  # 显式更新PRNG state
    return scale * jax.random.normal(w_key, (n, m)), scale * jax.random.normal(b_key, (n,))


def init_network_params(sizes, key):
    """Initialize all layers for a fully-connected neural network with sizes "sizes"
    """
    keys = jax.random.split(key, len(sizes))  # split可以同时创建多个key
    return [random_layer_params(m, n, k) for m, n, k in zip(sizes[:-1], sizes[1:], keys)]


layer_sizes = [784, 512, 512, 10]

key, init_key = jax.random.split(key)  # init_key used for initialization
params = init_network_params(layer_sizes, init_key)

print(len(params), len(params[0]), len(params[1]), len(params[2]))
# 3, 2, 2, 2

print(params[0][0].shape, params[0][1].shape)
# (512, 784), (512,)


# 创建网络，实际上就是写出forward
def relu(x):
    return jnp.maximum(0, x)

# 注意下面的x只是一张图片，我们并不需要自己动手去实现batched_x
def model_forward(params, x):
    # per-example predictions
    for w, b in params[:-1]:
        x = jnp.dot(w, x) + b
        x = relu(x)
    final_w, final_b = params[-1]
    logits = jnp.dot(final_w, x) + final_b
    return logits


# 模型forward已经完成了，下面测试下
key, test_key = jax.random.split(key)
random_flattened_image = jax.random.normal(test_key, (784, ))
preds = model_forward(params, random_flattened_image)
print(preds.shape)
# (10,)

# 创建一个随机batch数据, shape=(32, 784)
random_batched_flattened_images = jax.random.normal(jax.random.PRNGKey(1), (32, 784))
model_forward(params, random_batched_flattened_images)  # error
# TypeError: Incompatible shapes for dot: got (512, 784) and (32, 784).


# 创建支持batch数据的model_forward, 使用vmap so easy
batched_forward = vmap(model_forward, in_axes=(None, 0), out_axes=0)

batched_preds = batched_forward(params, random_batched_flattened_images)
print(batched_preds.shape)
# (32, 10)

import numpy as np
from torch.utils.data import DataLoader
from torchvision.datasets import MNIST
from torch.utils.data import Sampler, SequentialSampler


class FlattenAndCast(object):
    def __call__(self, pic):
        return np.ravel(np.array(pic, dtype=jnp.float32))

# DataLoader返回numpy array，而不是torch Tensor
def numpy_collate(batch):
    if isinstance(batch[0], np.ndarray):
        return np.stack(batch)
    elif isinstance(batch[0], (tuple,list)):
        transposed = zip(*batch)
        return [numpy_collate(samples) for samples in transposed]
    else:
        return np.array(batch)

class JAXRandomSampler(Sampler):
    def __init__(self, data_source, rng_key):
        self.data_source = data_source
        self.rng_key = rng_key
        
    def __len__(self):
        return len(self.data_source)
    
    def __iter__(self):
        self.rng_key, current_rng = jax.random.split(self.rng_key)
        return iter(jax.random.permutation(current_rng, jnp.arange(len(self))).tolist())


class NumpyLoader(DataLoader):
    def __init__(self, dataset, rng_key=None, batch_size=1,
                 shuffle=False, **kwargs):
        if shuffle:
            sampler = JAXRandomSampler(dataset, rng_key)
        else:
            sampler = SequentialSampler(dataset)
        
        super().__init__(dataset, batch_size, sampler=sampler, **kwargs)


# 借助于torchvision和NumpyLoader
mnist_dataset_train = MNIST('/tmp/mnist/', download=True, transform=FlattenAndCast())
key, loader_key = jax.random.split(key)
train_loader = NumpyLoader(mnist_dataset_train, loader_key, batch_size=32, shuffle=True,
                           num_workers=0, collate_fn=numpy_collate, drop_last=True)

mnist_dataset_test = MNIST('/tmp/mnist/', download=True, train=False, transform=FlattenAndCast())
eval_loader = NumpyLoader(mnist_dataset_test, batch_size=128, shuffle=False, num_workers=0,
                          collate_fn=numpy_collate, drop_last=False)
### Here we set num_workers=0

from jax.scipy.special import logsumexp


def loss(params, images, targets):
    logits = batched_forward(params, images)
    preds = logits - logsumexp(logits)
    return -jnp.mean(preds * targets)

@jit
def sgd_update(params, x, y, lr):
    grads = grad(loss)(params, x, y)
    return [(w - lr * dw, b - lr * db)
            for (w, b), (dw, db) in zip(params, grads)]

@jit
def sgd_update(params, x, y, lr):
    grads = grad(loss)(params, x, y)
    return [(param[0] - lr * grad[0], param[1] - lr * grad[1]) if len(param) == 2 else 
            (param[0] - lr * grad[0], param[1] - lr * grad[1], param[2] - lr * grad[2])
                for param, grad in zip(params, grads)]

[1, "a", object()]  # 这个list属于pytree，它含有3个leaf: 1, "a", object()

(1, (2, 3), ())  # 这个tuple属于pytree，含有三个leaf: 1, 2, 3

[1, {"k1": 2, "k2": (3, 4)}, 5]  # 这个list含有5个leaf: 1, 2, 3, 4, 5

ptree = (1, (2, 3), (), [(2,3,4,), 5], {"key": 2})
jax.tree_util.tree_map(lambda x: x+ 2, ptree)
# (3, (4, 5), (), [(4, 5, 6), 7], {'key': 4})

@jit
def sgd_update(params, x, y, lr):
    grads = grad(loss)(params, x, y)
    return jax.tree_util.tree_map(lambda p, g: p - lr * g, params, grads)
``

`请把优雅打在公屏上，栓Q`

# 训练流程和验证流程

ok，接下来就是把代码串起来，整一个训练流程和验证流程：


```python
def one_hot(x, k=10, dtype=jnp.float32):
    """Create a one-hot encoding of x of size k."""
    return jnp.array(x[:, None] == jnp.arange(k), dtype)

def accuracy(params, loader):
    total_acc = 0
    total_num = 0
    for x, y in loader:
        predicted_class = jnp.argmax(batched_forward(params, x), axis=1)
        total_num += len(x)
        total_acc += jnp.sum(predicted_class == y)
    return total_acc / total_num


lr = 0.01
n_classes = 10
for epoch in range(5):
    for idx, (x, y) in enumerate(train_loader):
        y = one_hot(y, n_classes)
        params = sgd_update(params, x, y, lr)
        lr = lr * 0.999 if lr > 1e-3 else 1e-3  # very simple lr scheduler
        if idx % 20 == 0:  # evaluation
            train_acc = accuracy(params, train_loader)
            eval_acc = accuracy(params, eval_loader)
            print("Epoch {} - batch_idx {}, Training set acc {}, eval set accuracy {}".format(
                  epoch, idx, train_acc, eval_acc))

import jax
from jax import numpy as jnp
from jax import grad, jit, vmap
from flax import linen as nn
from typing import Sequence

# 创建 PRNGKey (PRNG State)
key = jax.random.PRNGKey(0)


class MLP(nn.Module):
    layer_sizes: Sequence[int] = None  # 类型标注信息 Sequence[int]
    
    def setup(self):
        # 创建Dense()时只设置了输出维度大小，输入维度大小需要Flax进行推测
        self.layers = [nn.Dense(features=size) for size in self.layer_sizes[1:]]
    
    def __call__(self, x):
        for layer in self.layers[:-1]:
            x = layer(x)
            x = nn.relu(x)
        return self.layers[-1](x)

layer_sizes = [784, 512, 512, 10]


# 创建模型
model = MLP(layer_sizes)

# 使用`init`和dummy_x来创建模型参数
# 注意，在创建Dense()时并没有指定输入维度大小，
# `init`方法本质上也是调用`__call__`，利用dummpy_x来进行推测参数的shape
key, init_key = jax.random.split(key)  # init_key used for initialization
dummy_x = jax.random.uniform(init_key, (784, ))
key, init_key = jax.random.split(key)

# init_key
params = model.init(init_key, dummy_x)

# params

# 创建一个随机batch数据, shape=(32, 784)
# 模型同Pytorch Module一样，自动支持batch数据，所以也不需要手动vmap了
random_batched_flattened_images = jax.random.normal(jax.random.PRNGKey(1), (32, 784))
model.apply(params, random_batched_flattened_images).shape
# (32, 10)

import optax

lr = 1e-3

# 学习率调度算法
lr_decay_fn = optax.linear_schedule(
        init_value=lr,
        end_value=1e-5,
        transition_steps=200,
)

# 直接上Adam
optimizer = optax.adam(
            learning_rate=lr_decay_fn,
)

state = train_state.TrainState.create(apply_fn=model.apply, params=params, tx=optimizer)

def train_step(state, x, y):
    """Computes gradients and loss for a single batch."""
    def loss_fn(params):
        logits = state.apply_fn(params, x)
        one_hot = jax.nn.one_hot(y, 10)
        loss = jnp.mean(optax.softmax_cross_entropy(logits=logits, labels=one_hot))
        return loss

    grad_fn = value_and_grad(loss_fn)
    loss, grads = grad_fn(state.params)
    new_state = state.apply_gradients(grads=grads)
    return new_state, loss

# donate_argnums用于buffer复用，这里指的是输入和输出的state buffer复用
jit_train_step = jit(train_step, donate_argnums=(0,))  


@jax.jit
def apply_model(state, x):
    """Computes gradients and loss for a single batch."""
    
    logits = state.apply_fn(state.params, x)
    return jnp.argmax(logits, -1)

def eval_model(state, loader):
    total_acc = 0.
    total_num = 0.
    for x, y in loader:
        y_pred = apply_model(state, x)
        total_num += len(x)
        total_acc += jnp.sum(y_pred == y)
    return total_acc / total_num

for epoch in range(5):
    for idx, (x, y) in enumerate(train_loader):
        state, loss = jit_train_step(state, x, y)
        if idx % 20 == 0:  # evaluation
            train_acc = eval_model(state, train_loader)
            eval_acc = eval_model(state, eval_loader)
            print("Epoch {} - batch_idx {}, loss {}, Training set acc {}, eval set accuracy {}".format(
              epoch, idx, loss, train_acc, eval_acc))

class MLP(nn.Module):
    
    def setup(self):
        self.layer1 = nn.Dense(features=512)
        self.dropout1 = nn.Dropout(rate=0.3)
        self.norm1 = nn.BatchNorm()
        
        self.layer2 = nn.Dense(features=512)
        self.dropout2 = nn.Dropout(rate=0.4)
        self.norm2 = nn.BatchNorm()
        
        self.layer3 = nn.Dense(features=10)
        
    def __call__(self, x, train:bool = True):
        """train用于区分train_mode or eval_mode"""

        x = nn.relu(self.layer1(x))
        x = self.dropout1(x, deterministic=not train)
        x = self.norm1(x, use_running_average=not train)
        x = nn.relu(self.layer2(x))
        x = self.dropout2(x, deterministic=not train)
        x = self.norm2(x, use_running_average=not train)
        
        x = self.layer3(x)

        return x

# 创建模型
model = MLP()

# 使用`init`和dummy_x来创建模型参数
key, init_key = jax.random.split(key)
dummy_x = jax.random.uniform(init_key, (784, ))

key, init_key, drop_key = jax.random.split(key, 3)  # 通过split得到3个key
# "dropout"这个名字是固定的
variables = model.init({"params": init_key, "dropout": drop_key}, dummy_x, train=True)

variables.keys()
# frozen_dict_keys(['params', 'batch_stats'])

variables['batch_stats'].keys()
# frozen_dict_keys(['norm1', 'norm2'])

variables['batch_stats']['norm1'].keys()
# frozen_dict_keys(['mean', 'var'])

key, drop_key = jax.random.split(key)

y, non_trainable_params = model.apply(variables, dummy_x, train=True, rngs={"dropout": drop_key},
                                      mutable=['batch_stats']) 

non_trainable_params.keys()
# frozen_dict_keys(['batch_stats'])

class CustomTrainState(train_state.TrainState):
    batch_stats: flax.core.FrozenDict[str, Any]


state = CustomTrainState.create(
    apply_fn=model.apply,
    params=variables['params'],
    tx=optimizer,
    batch_stats=variables['batch_stats'],
)

def train_step(state, x, y, dropout_key):
    """Computes gradients and loss for a single batch."""
    def loss_fn(params):
        logits, new_state = state.apply_fn({"params": params, "batch_stats": state.batch_stats},
                                           x, train=True, rngs={"dropout": dropout_key}, mutable=["batch_stats"])
        
        one_hot = jax.nn.one_hot(y, 10)
        loss = jnp.mean(optax.softmax_cross_entropy(logits=logits, labels=one_hot))
        return loss, new_state

    grad_fn = value_and_grad(loss_fn, has_aux=True)  # `value_and_grad`在进行grad同时返回loss
    (loss, new_state), grads = grad_fn(state.params)
    new_state = state.apply_gradients(grads=grads, batch_stats=new_state["batch_stats"])
    return new_state, loss

jit_train_step = jit(train_step, donate_argnums=(0,))  # donate_argnums用于buffer复用，这里指的是输入和输出的state buffer复用


@jax.jit
def apply_model(state, x):
    """Computes gradients and loss for a single batch."""
    
    logits = state.apply_fn({"params":state.params, "batch_stats": state.batch_stats},
                            x, train=False)  # train设置为False，即为eval mode
    return jnp.argmax(logits, -1)

for epoch in range(5):
    for idx, (x, y) in enumerate(train_loader):
        key, dropout_key = jax.random.split(key)
        state, loss = jit_train_step(state, x, y, dropout_key)
        if idx % 100 == 0:  # evaluation
            train_acc = eval_model(state, train_loader)
            eval_acc = eval_model(state, eval_loader)
            print("Epoch {} - batch_idx {}, loss {}, Training set acc {}, eval set accuracy {}".format(
              epoch, idx, loss, train_acc, eval_acc))