Seungwoo (Simon) Kim

> It seems there is a device placement issue with torch on GPU, could you take a look? https://btx.cloud.google.com/invocations/0b42aedc-c253-401a-b110-1f441959ccc8/targets/keras%2Fgithub%2Fubuntu%2Fgpu%2Ftorch%2Fcontinuous/log Thanks for the info! I made a PR with a fix:...

> @SeKim12 thanks for this nice contribution! > > Did you test using the pytorch backend? Are you able to compile the associative scan operator? I will be testing this...

Hi @SamanehSaadat, My apologies -- I just saw this notification! It seems like the optree imports were handled by @fchollet. Thank you! Please let me know if there is anything...

From what I could tell, it's supported in [tensorflow](https://www.tensorflow.org/probability/api_docs/python/tfp/math/scan_associative), jax (above), but not yet in PyTorch or Numpy. I'd be interested in working on this op. I also see a...

@edgarriba @ducha-aiki Hi, please let me know if there is anything else I can work on for this PR. I don't think the failing checks come from this PR. Thanks!