Jax框架的static与Traced Operations —— Static vs Traced Operations

相关：

Jax框架的jit编译是否可以使用循环结构，如果使用循环结构需要注意什么

---

Jax的static和Traced都是指jit编译的函数内的对象的属性的，jit装饰的函数其输入参数和输出参数都是Traced类型的（除了使用@partial装饰符提前指定为static的），而jit装饰的函数内的所有jax的array都是Traced的，即使不依赖输入的参数也是为Traced的，如：

import jax.numpy as jnp
from jax import jit

@jit
def f(x):
  # return x.reshape(jnp.array(x.shape).prod())
  print(x)
  y = jnp.ones((3,3))
  print(y)

x = jnp.ones((2, 3))
f(x)

但是，jit装饰的函数内的所有python类型的array都是static的，如：

也就是说，jit在编译时会把python的代码直接进行编译，硬编码到后端的代码中，这样所有在jit内的python代码都为static类型。

所有的jit包装下的函数中的循环和判断语句都不能依赖于Traced对象，也就是说jit包装的函数在编译时需要将判断语句取出固定的分支，对循环语句进行展开，可以把jit下的循环和判断语句视作static。

---

static类型，是指在jit编译时就把具体的数值编译到代码中，把循环和判断语句全部转为顺序结构的语句。

Traced类型，是指在jit编译时将变量的shape和type固定，jit编译后的代码可以通过Tracd对象的不同值来得出不同的运算结果，但是不同的运行次中Traced对象的shape和type都是不能变的，如果调用过程中Traced对象的shape或type改变那么jax的jit会对代码进行重新编译并缓存。

---

在jit的包装中，jax的对象（因为是jax的类型对象必然是Traced类型）的一些操作，要求操作的参数必须是static类型，也就是说jit编译时有些jax的一些操作的参数要求在编译时是需要确定的，如jnp.array对象的reshape操作。

因为Traced对象的shape和type在jit编译时需要固定的，也就是说Traced对象的shape和type是static类型的，因此对jnp.array的static类型的shape进行reshape操作也需要操作参数为static的，即jnp.array的reshape参数和该jnp.array的shape参数都需要在jit编译时指定。

官方给出的一个例子：

由于jnp.array(x.shape).prod()操作得到的变量依旧为jnp.array类型，即jax类型，那么也就为Traced类型，但是jnp.array的reshape操作要求参数为static类型，即在jit编译时硬编码到代码中，因此报错。

关于这一点，我们要注意，对于jax的一些对象的操作，会要求操作参数或操作对象为static或Traced的要进行区分。

对于jit下的一些static类型变量通过numpy或python操作后获得的新变量依旧为static类型。而不论是对static对象还是Traced对象进行jax操作那么就会得到Traced类型对象。

对static类型对象进行numpy或python操作，依旧获得python或numpy类型的static对象：

对Traced对象进行操作，得到的依旧为Traced类型：

---

对于前面最早给出的报错的例子，下面给出修改后的代码：

from jax import jit
import jax.numpy as jnp
import numpy as np

@jit
def f(x):
  return x.reshape((np.prod(x.shape),))

f(x)