AdamW State Bytes

Problem

Compute the AdamW optimizer state size in bytes. AdamW keeps two moment estimates per parameter, so the state is 2 * num_params * bytes_per_value.

The signature is adamw_state_bytes(num_params, bytes_per_value=4). Return an integer.

Hint

The two moments are the first moment and the second moment.

case 1

in: [1,4]

expected: 8

case 2

in: [1000000,4]

expected: 8000000

case 3

in: [1000,2]

expected: 4000

case 4

in: [0,4]

expected: 0

Starter codePython

def adamw_state_bytes(num_params, bytes_per_value=4):
    # Your code here
    pass

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