Allan Zhou

Error running particle experiments

2

Hi, thanks for open sourcing this work! I tried running: ``` ./ibc/ibc/configs/particle/run_mlp_ebm_langevin_best.sh 2 ``` And got this error ``` File "ibc/ibc/train_eval.py", line 397, in main [122/528] strategy=strategy) File "/iris/u/ayz/anaconda3/envs/ibc/lib/python3.7/site-packages/gin/config.py", line...

Pallas+scan: NotImplementedError when num_extensive is True

1

### Description I tried to write a simple TPU pallas kernel to implement cum_sum, where the input is being chunked along the summed dimension and the kernel is calculating one...