Wasm: implementing cumsum

[davidpham87]

Hello,

I would be interested to have a cumulative sum/prod in WASM. The reason is I want to implement a rolling average.

I wonder if the reductions (see Clojure reductions) would make sense.

Best regards,
David

[alvinsunyixiao]

Is there any recent plans for supporting this for the wasm backend?

[rthadur]

It is in our todo list , there is no fixed timeline as of now , please let us know if you would like to contribute.

[alvinsunyixiao]

@rthadur I would like to give it a try. But I have limited TypeScript experience, is there a starting point or examples of other operators I could follow?

[annxingyuan]

Hi @alvinsunyixiao - WASM cumsum would be an awesome contribution! Here is a recent PR adding support for reverse to WASM: #3526

I think cumsum could follow a similar format. You can use the CPU implementation as a reference for the kernel logic: https://github.com/tensorflow/tfjs/blob/master/tfjs-backend-cpu/src/backend_cpu.ts#L619

In case it's helpful, here's a guide to contributing: https://github.com/tensorflow/tfjs/blob/master/CONTRIBUTING.md

Please let us know if you run into any snags!

[alvinsunyixiao]

Hi @annxingyuan, thanks for the pointer! I just made time for starting to implement this. I noticed the CPU backend implementation asserted axis === x.rank - 1. Do we expect the same behavior for wasm? Or should I try generalize this to all axis?

[annxingyuan]

Hi @alvinsunyixiao - I would say maybe try to implement it for the axis === x.rank - 1 case first, then we can generalize it in a follow-up PR!

[alvinsunyixiao]

Hi @annxingyuan - What are the supported dtypes cumsum should implement? I noticed the Reverse kernel assumes a float32 type, but since that is just a copy of some sort, the code could also work for int32 because of the memory pointer size being the same. However, for cumsum, should I write additional logic to support int32 tensors?

[alvinsunyixiao]

It seems like Sum is also assuming a float32 input, but I didn't see an assertion anywhere ensuring that dtype. Is that expected?

[alvinsunyixiao]

Hi @annxingyuan - I have just created a PR #3782. Thanks for the help!

[rthadur]

related PR is merged now , closing this issue.