Tensorflow XLANG位发生器占用太多内存

XLA为这个张量分配4G内存。其大小似乎与批量大小成比例。这对我来说毫无意义，它似乎不是要存储在HBM中的模型图的一部分。我用的是TPUv3

除了初始化模型外，我不使用任何随机操作。此外，我为所有权重声明了bfloat16，但这是一个u32张量

  Largest program allocations in hbm:

  1. Size: 4.00G
     Shape: u32[128,8,1024,1024]{3,2,1,0:T(8,128)}
     Unpadded size: 4.00G
     XLA label: %rng-bit-generator = (u32[2,128]{1,0:T(2,128)}, u32[128,8,1024,1024]{3,2,1,0:T(8,128)}) rng-bit-generator(u32[2,128]{1,0:T(2,128)} %fusion.2446), algorithm=rng_default
     Allocation type: HLO temp
     ==========================

上述分配的原因可能是什么？我使用的像素来自：

问题：

• 当我所有的权重/模型定义（包括全局环境标志）都使用BF16时，为什么该张量的类型为

  u32

• 为什么使用rng位发生器

---

看起来合理：4 GB=128*8*1024*1024*4@Andrey毫无疑问，这是正确的。但是，我找不到它是否对应于我的模型的任何组件。此外，rng位发生器更不清楚为什么需要它。它还使用u32，而不是我指定的bfloat16。