Question regarding cpu state representation

[anthony0727]

We're trying to reproduce your work, I was wondering how you encoded register and memory.

Can you tell me what each of dimension stands for? We assume B for batch size, and one for array of 32bits, one for number of inputs, and one for registers and memories(e.g. rax, rbx, ..., stack, heap, ...) https://github.com/deepmind/alphadev/blob/709dfb1b36f8efa6134145f0563bfb4cb3f8f923/alphadev.py#L481

Additionally num_location_values seems to be undefined variable, shouldn't it be num_locations from TaskSpec? https://github.com/deepmind/alphadev/blob/709dfb1b36f8efa6134145f0563bfb4cb3f8f923/alphadev.py#L487