MediumCalculusca-206

Checkpoint Chain

Problem

Gradient checkpointing stores activations every k of n layers instead of all of them, trading memory for recomputation.

The number of checkpoints is ceil(n / k), and backward recomputation adds one extra layer evaluation per checkpoint, so total forward-equivalent passes = n + ceil(n/k). Given n and k, return [checkpoints, total_passes].

Hint

Smaller k saves more memory but costs more recomputation.

Test cases

case 1

in: [6,2]

expected: [3,9]

case 2

in: [10,3]

expected: [4,14]

case 3

in: [1,1]

expected: [1,2]

case 4

in: [7,4]

expected: [2,9]

Starter codePython

def checkpoint_chain(n, k):
    # Return [checkpoints, total_passes]
    # Your code here
    pass

4 test cases · run your code against them in the editor