Implement:
mlp_params(sizes, bias=True): parameter count of an MLP with layer widths sizes, e.g. [784, 256, 10] has two linear layerstraining_memory_bytes(n_params, bytes_per_param=4, optimizer="adam"): memory for weights + gradients + optimizer state, where Adam keeps 2 extra values per parameter, SGD with momentum 1, and plain SGD 0