Fix command buffer memory tracking to use bytes instead of elements

[hxu296]

Proposed changes

buffer_sizes is compared against a megabyte threshold (can be set by user with MLX_MAX_MB_PER_BUFFER), but was accumulating data_size() which is in elements, not bytes. This made the commit heuristic dtype-blind, a 4MB float32 buffer counted the same as a 2MB float16 buffer with the same number of elements.

Checklist

Put an x in the boxes that apply.

• I have read the CONTRIBUTING document
• I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
• I have added tests that prove my fix is effective or that my feature works (please lmk if new tests are needed)
• I have updated the necessary documentation (if needed)

[angeloskath]

Yep not necessarily a bad change. We 'll have to do some perf testing before merging this though (and adjust the heuristics). Given that this is a heuristic for the time an op will take it isn't necessarily bad that it is in elements, so if it does work better we might have to change the env variable name.

[hxu296]

Makes sense. The change will make command buffer commits more frequent for larger dtypes, so small workloads may not fully hide command buffer encode/schedule overhead. Definitely worth benchmarking and potentially tuning the thresholds / just change env var name for now.