Optimizer State Memory

Problem

Estimate optimizer state memory: num_params * num_states * bytes_per_value.

The signature is optimizer_state_memory(num_params, num_states, bytes_per_value=4). Return an integer number of bytes.

Hint

Adam-style optimizers keep two state tensors per parameter.

case 1

in: [1000000,2,4]

expected: 8000000

case 2

in: [1,2,4]

expected: 8

case 3

in: [7000000000,2,4]

expected: 56000000000

case 4

in: [1000,1,2]

expected: 2000

Starter codePython

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

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