Adding some findings so far:

1. the difference between cummax/cummin and cumprod/cumsum is that the pytorch API also returns a tensor representing the indices of the max values in the tensor: https://pytorch.org/docs/stable/generated/torch.cummax.html.
2. The current implementation of cumsum uses the xla::ReduceWindowWithGeneralPadding function with the xla:Add computation.
3. However, if we were to get the indices of the min/max values, we can probably concatenate the original tensor with iota as indicies and use a computation that keeps both the values and indices. Need to investigate more to see if this is already supported somehow.

Might be able to follow argmax here: https://github.com/openxla/xla/blob/d09a2efca944d1f5ae668006111be81bc28750bd/xla/hlo/builder/lib/arithmetic.cc#L144-L172