Jianing Ye

Stacking layers with `vmap` and forwarding with `scan` result in loss of precision in XLA backend

2

I followed the instructions in the [tutorial#scan-over-layers](https://flax.readthedocs.io/en/latest/nnx_basics.html#scan-over-layers) to build a network with multiple layers with `nnx.vmap` and to forward with `nnx.scan`. However, doing so reults in loss of precision in...

[BUG] Replay buffer on CPU is much slower than the performance in doc.

12

### Describe the bug I am trying to put the buffer on CPU to alleviate the memory shortage on GPU, but encountered a large speed degradation. As shown in `readme.md`,...

`loop_fn` in `search.py` becomes slow as the tree depth increases

1

I am using mctx to implement MuZero, where actions are selected at each node via a forward pass of a large neural network. Initially, the main bottleneck in terms of...