jax.random.randint在spu环境下结果似乎不能正确运行

[Maxime-Tong]

Issue Type

Bug

Source

binary

Secretflow Version

0.7.18b0

OS Platform and Distribution

Linux centos 3.10.0

Python version

3.8.13

Bazel version

No response

GCC/Compiler version

No response

What happend and What you expected to happen.

根据jax文档，jax.random.randint函数可以根据key生成一个可以选定范围大小的数字，但是在spu环境下，这个范围大小似乎不成立。
如下面给出的代码，尝试循环100个epoch，生成100个在[0,5)范围的随机数，统计超出范围的随机数个数。
在jit环境下，返回的统计个数为0；而在spu环境下，返回的统计个数却是100。这是为什么？

Reproduction code to reproduce the issue.

def run_on_cpu():
    def train():
        def loop_fun(_, carry):
            key, cnt = carry
            key, subkey = random.split(key)
            j = random.randint(subkey, shape=(), minval=0, maxval=5, dtype=jnp.int32)
            delta = jnp.where(jnp.any(jnp.array([j>=5, j<0])), 1, 0)
            cnt += delta
            return (key, cnt)
        key = random.PRNGKey(42)
        return lax.fori_loop(0, 100, loop_fun, (key, 0))
    key, cnt = jax.jit(train)()
    print(key, cnt)
# > [3622954601 1653733208] 0

# spu环境下
def run_on_spu():
    @ppd.device("SPU")
    def train():        
        def loop_fun(_, carry):
            key, cnt = carry
            key, subkey = random.split(key)
            j = random.randint(subkey, shape=(), minval=0, maxval=5, dtype=jnp.int32)
            delta = jnp.where(jnp.any(jnp.asarray([j>=5, j<0])), 1, 0)
            cnt += delta
            return (key, cnt)
        key = random.PRNGKey(42)
        key, cnt = lax.fori_loop(0, 100, loop_fun, (key, 0))
        return key, cnt
     key, cnt = train()
     key, cnt = ppd.get(key), ppd.get(cnt)
     print(key, cnt)
#> [3463333171 2461478805] 100

[anakinxc]

Hi @Maxime-Tong

jax.random 的支持现在的确是有点问题。我来大概解释一下，jax 在生成随机数的过程里用了大量连续的位操作和运算，SPU 因为是用的一个定长的 ring 来模拟所有的整型和浮点数，在这个过程中一旦实际行为和 jax 的预期出现偏差就会导致随机数完全不在范围内

目前的建议是随机数可以在明文里生成然后传给 spu。您提供的这个 repro step 我们会内部分析一下，然后争取在后面的跌代里修复这个问题

祝新年快乐 :P

[Maxime-Tong]

好的，谢谢，另外似乎jax提供的三种控制流的方法中，lax.cond和lax.while_loop函数在spu环境下还不能正常运行，这是正常的吗？

另同祝新年快乐 : )

[anakinxc]

好的，谢谢，另外似乎jax提供的三种控制流的方法中，lax.cond和lax.while_loop函数在spu环境下还不能正常运行，这是正常的吗？

另同祝新年快乐 : )

cond 不应该，while_loop 有个限制是 condition 结果必须是明文的

最好可以提供一个有问题的 example，我来看看