Corrected some bugs in FlatSum and added a DifferentiableFlatSum that accumulates the gradients