Jax中关于NonZero的使用
技术背景
在Jax的JIT即时编译中，会追踪每一个Tensor的Shape变化。如果在计算的过程中出现了一些动态Shape的Tensor（Shape大小跟输入的数据有关），那么就无法使用Jax的JIT进行编译优化。最常见的就是
numpy.where
这种操作，因为这个操作返回的是符合判定条件的Index序号，而不同输入对应的输出Index长度一般是不一致的，因此在Jax的JIT中无法对该操作进行编译。当然，需要特别说明的是，
numpy.where
这个操作有两种用法，一种是
numpy.where(condition, 1, 0)
直接给输入打上Mask。另一种用法是
numpy.where(condition)
，这种用法返回的就是一个Index序列，也就是我们需要讨论的应用场景。
普通模式
我们考虑一个比较简单的Toy Model用于测试：
\]
在不采用即时编译的场景下，Jax的代码可以这么写：
import os
os.environ['XLA_PYTHON_CLIENT_PREALLOCATE']='false'
import numpy as np
np.random.seed(0)
import jax
from jax import numpy as jnp
def func(r, q, cutoff=0.2):
dis = jnp.abs(r[:, None] - r[None])
maski, maskj = jnp.where(dis<=cutoff)
qi = q[maski]
qj = q[maskj]
return jnp.sum(qi*qj)
N = 100
r = jnp.array(np.random.random(N), jnp.float32)
q = jnp.array(np.random.random(N), jnp.float32)
print (func(r, q))
# 1035.7422
那么我们先记住这个输出的结果，因为采用的随机种子是一致的，一会儿可以直接跟JIT的输出结果进行对比。
JIT模式
Jax的JIT模式的使用方法常见的就是三种，一种是在函数头顶加一个
装饰器
，一种是在函数引用的时候使用
jax.jit(function)
来调用，最后一种是配合
partial
偏函数
来使用，都不是很复杂。那么这里先用装饰器的形式演示一下Jax中即时编译的用法：
import os
os.environ['XLA_PYTHON_CLIENT_PREALLOCATE']='false'
import numpy as np
np.random.seed(0)
import jax
from jax import numpy as jnp
@jax.jit
def func(r, q, cutoff=0.2):
dis = jnp.abs(r[:, None] - r[None])
maski, maskj = jnp.where(dis<=cutoff)
qi = q[maski]
qj = q[maskj]
return jnp.sum(qi*qj)
N = 100
r = jnp.array(np.random.random(N), jnp.float32)
q = jnp.array(np.random.random(N), jnp.float32)
print (func(r, q))
正如前面所说，因为
numpy.where
对应的输出是一个动态的Shape，那么在编译阶段就会报错。报错信息如下：
Traceback (most recent call last):
File "/home/dechin/projects/gitee/dechin/tests/jax_mask.py", line 21, in <module>
print (func(r, q))
File "/home/dechin/anaconda3/envs/jax/lib/python3.10/site-packages/jax/_src/traceback_util.py", line 162, in reraise_with_filtered_traceback
return fun(*args, **kwargs)
File "/home/dechin/anaconda3/envs/jax/lib/python3.10/site-packages/jax/_src/api.py", line 622, in cache_miss
execute = dispatch._xla_call_impl_lazy(fun_, *tracers, **params)
File "/home/dechin/anaconda3/envs/jax/lib/python3.10/site-packages/jax/_src/dispatch.py", line 236, in _xla_call_impl_lazy
return xla_callable(fun, device, backend, name, donated_invars, keep_unused,
File "/home/dechin/anaconda3/envs/jax/lib/python3.10/site-packages/jax/linear_util.py", line 303, in memoized_fun
ans = call(fun, *args)
File "/home/dechin/anaconda3/envs/jax/lib/python3.10/site-packages/jax/_src/dispatch.py", line 359, in _xla_callable_uncached
return lower_xla_callable(fun, device, backend, name, donated_invars, False,
File "/home/dechin/anaconda3/envs/jax/lib/python3.10/site-packages/jax/_src/profiler.py", line 314, in wrapper
return func(*args, **kwargs)
File "/home/dechin/anaconda3/envs/jax/lib/python3.10/site-packages/jax/_src/dispatch.py", line 445, in lower_xla_callable
jaxpr, out_type, consts = pe.trace_to_jaxpr_final2(
File "/home/dechin/anaconda3/envs/jax/lib/python3.10/site-packages/jax/_src/profiler.py", line 314, in wrapper
return func(*args, **kwargs)
File "/home/dechin/anaconda3/envs/jax/lib/python3.10/site-packages/jax/interpreters/partial_eval.py", line 2077, in trace_to_jaxpr_final2
jaxpr, out_type, consts = trace_to_subjaxpr_dynamic2(fun, main, debug_info)
File "/home/dechin/anaconda3/envs/jax/lib/python3.10/site-packages/jax/interpreters/partial_eval.py", line 2027, in trace_to_subjaxpr_dynamic2
ans = fun.call_wrapped(*in_tracers_)
File "/home/dechin/anaconda3/envs/jax/lib/python3.10/site-packages/jax/linear_util.py", line 167, in call_wrapped
ans = self.f(*args, **dict(self.params, **kwargs))
File "/home/dechin/projects/gitee/dechin/tests/jax_mask.py", line 12, in func
maski, maskj = jnp.where(dis<=cutoff)
File "/home/dechin/anaconda3/envs/jax/lib/python3.10/site-packages/jax/_src/numpy/lax_numpy.py", line 1077, in where
return nonzero(condition, size=size, fill_value=fill_value)
File "/home/dechin/anaconda3/envs/jax/lib/python3.10/site-packages/jax/_src/numpy/lax_numpy.py", line 1332, in nonzero
size = core.concrete_or_error(operator.index, size,
File "/home/dechin/anaconda3/envs/jax/lib/python3.10/site-packages/jax/core.py", line 1278, in concrete_or_error
raise ConcretizationTypeError(val, context)
jax._src.traceback_util.UnfilteredStackTrace: jax._src.errors.ConcretizationTypeError: Abstract tracer value encountered where concrete value is expected: Traced<ShapedArray(int32[])>with<DynamicJaxprTrace(level=0/1)>
The size argument of jnp.nonzero must be statically specified to use jnp.nonzero within JAX transformations.
The error occurred while tracing the function func at /home/dechin/projects/gitee/dechin/tests/jax_mask.py:9 for jit. This concrete value was not available in Python because it depends on the value of the argument 'r'.
See https://jax.readthedocs.io/en/latest/errors.html#jax.errors.ConcretizationTypeError
The stack trace below excludes JAX-internal frames.
The preceding is the original exception that occurred, unmodified.
--------------------
The above exception was the direct cause of the following exception:
Traceback (most recent call last):
File "/home/dechin/projects/gitee/dechin/tests/jax_mask.py", line 21, in <module>
print (func(r, q))
File "/home/dechin/projects/gitee/dechin/tests/jax_mask.py", line 12, in func
maski, maskj = jnp.where(dis<=cutoff)
File "/home/dechin/anaconda3/envs/jax/lib/python3.10/site-packages/jax/_src/numpy/lax_numpy.py", line 1077, in where
return nonzero(condition, size=size, fill_value=fill_value)
File "/home/dechin/anaconda3/envs/jax/lib/python3.10/site-packages/jax/_src/numpy/lax_numpy.py", line 1332, in nonzero
size = core.concrete_or_error(operator.index, size,
jax._src.errors.ConcretizationTypeError: Abstract tracer value encountered where concrete value is expected: Traced<ShapedArray(int32[])>with<DynamicJaxprTrace(level=0/1)>
The size argument of jnp.nonzero must be statically specified to use jnp.nonzero within JAX transformations.
The error occurred while tracing the function func at /home/dechin/projects/gitee/dechin/tests/jax_mask.py:9 for jit. This concrete value was not available in Python because it depends on the value of the argument 'r'.
See https://jax.readthedocs.io/en/latest/errors.html#jax.errors.ConcretizationTypeError
想避免这个报错，要么就是对该函数不做编译（牺牲性能），要么就是自己写一个CUDA算子（增加工作量），再就是我们这里用到的NonZero定长输出的方法（预置条件）。
NonZero的使用
使用Jax的NonZero函数时，也有一点需要注意：虽然NonZero可以做到固定长度的输出，但是这个固定的长度本身也是一个名为
size
的传入参数。也就是说，NonZero的输出Shape也是要取决于输入参数的。Jax开发时也考虑到了这一点，所以在编译时提供了一个功能可以设置静态参量：
static_argnames
，例如我们的案例中，将
size
这个名称的传参设置为静态参量，这样就可以使用Jax的即时编译了：
import os
os.environ['XLA_PYTHON_CLIENT_PREALLOCATE']='false'
import numpy as np
np.random.seed(0)
import jax
from jax import numpy as jnp
from functools import partial
@partial(jax.jit, static_argnames='size')
def func(r, q, cutoff=0.2, size=5000):
if q.shape[0] != r.shape[0]+1:
raise ValueError("The q.shape[0] should be equal to r.shape[0]+1")
dis = jnp.abs(r[:, None] - r[None])
maski, maskj = jnp.nonzero(jnp.where(dis<=cutoff, 1, 0), size=size, fill_value=-1)
qi = q[maski]
qj = q[maskj]
return jnp.sum(qi*qj)
N = 100
r = jnp.array(np.random.random(N), jnp.float32)
q = jnp.array(np.random.random(N), jnp.float32)
pader = jnp.array([0.], jnp.float32)
q = jnp.append(q, pader)
print (func(r, q))
# 1035.7422
可以看到，函数用Jax的JIT成功编译，并且输出结果跟前面未编译时候是一致的。当然，这里还用到了一个小技巧，就是NonZero函数输出结果时，不到长度的输出结果会被自动Pad到给定的长度，这里Pad的值使用的是我们给出的
fill_value
。因为NonZero输出的也是索引，这样我们可以把Pad的这些索引设置为
-1
，然后在构建参数
\(q\)
的时候事先在末尾
append
一个0，这样就可以确保计算的输出结果直接就是正确的。
总结概要
在使用Jax的过程中，有时候会遇到函数输出是一个动态的Shape，这种情况下我们很难利用到Jax的即时编译的功能，不能使得性能最大化。这也是使用Tensor数据结构来计算的一个特点，有好有坏。本文介绍了Jax的另外一个函数NonZero，可以使得我们能够编译那些动态Shape输出的函数。
版权声明
本文首发链接为：
https://www.cnblogs.com/dechinphy/p/nonzero.html
作者ID：DechinPhy
更多原著文章：
https://www.cnblogs.com/dechinphy/
请博主喝咖啡：
https://www.cnblogs.com/dechinphy/gallery/image/379634.html
