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用于测试：

\[E=\sum_{|r_i-r_j|\leq \epsilon}q_iq_j \]

在不采用即时编译的场景下，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)
padder = jnp.array([0.], jnp.float32)
q = jnp.append(q, padder)

print (func(r, q))
# 1035.7422

可以看到，函数用Jax的JIT成功编译，并且输出结果跟前面未编译时候是一致的。当然，这里还用到了一个小技巧，就是NonZero函数输出结果时，不到长度的输出结果会被自动Pad到给定的长度，这里Pad的值使用的是我们给出的fill_value。因为NonZero输出的也是索引，这样我们可以把Pad的这些索引设置为-1，然后在构建参数\(q\)的时候事先在末尾append一个0，这样就可以确保计算的输出结果直接就是正确的。