# Min-cut optimal(\*) recomputation (i.e. activation checkpointing) with AOTAutograd **URL:** \ **Category:** compiler\ **Created:** [January 10, 2022, 8:43pm UTC](https://dev-discuss.pytorch.org/t/min-cut-optimal-recomputation-i-e-activation-checkpointing-with-aotautograd/467 "2022-01-10T20:43:45Z")\ **Posts on this page:** 1\ **Showing post:** 8

**Author:** ![kelayamatoz](https://yyz2.discourse-cdn.com/flex036/user_avatar/dev-discuss.pytorch.org/kelayamatoz/32/938_2.png) [@kelayamatoz](https://dev-discuss.pytorch.org/u/kelayamatoz)\ **Post date:** [February 13, 2023, 7:46pm UTC](https://dev-discuss.pytorch.org/t/min-cut-optimal-recomputation-i-e-activation-checkpointing-with-aotautograd/467/8 "2023-02-13T19:46:04Z")

This is a very cool feature! Formulating this problem as a min-cut is quite neat! Quick question: Is there a way to configure recomputable\_ops based on their compute / data bandwidth ratio? It looks to me that the current recomputable\_ops list is hard-coded based on the types of operators, which could miss some optimization opportunities. For example, let’s say that we are looking at A @ B with two configurations: Config. 1: A.shape = (100, 100), B.shape = (100, 100), Dim\_M = Dim\_K = Dim\_N = 100 Config. 2: A.shape = (100, 2), B.shape = (2, 2), Dim\_M = 100, Dim\_K = 2, Dim\_N = 2 Config 1 and 2 are both MatMuls and should be unrecomputable\_ops. However, Config 2 makes compute / memory footprint of A @ B look more like an element-wise operator in the sense that ops per Dim\_M is very low (Dim\_K and Dim\_N \<\< Dim\_M). Therefore, Config 2 actually makes A @ B a good candidate to recompute. If we want to make Min-cut recomputation shape-aware, what in your view could be a good way to approach this? Thanks! --- _[View the full topic](https://dev-discuss.pytorch.org/t/min-cut-optimal-recomputation-i-e-activation-checkpointing-with-aotautograd/467)._