I am looking for a way to efficiently flatten gradients into a 1-D tensor, perform some operations, and then convert it back to the initial format in tf2. Does TensorFlow provid