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

发布时间 2024-01-10 00:15:42作者: Angry_Panda

相关:
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)

image


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

image

也就是说,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编译时指定。


官方给出的一个例子:

image

由于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对象:

image


对Traced对象进行操作,得到的依旧为Traced类型:

image


image