[New feature] Support exporting memory snapshots compatible with PyTorch's memory viz format

Hi, I'm currently implementing support for dumping memory snapshots that will be compatible with PyTorch's memory viz [web app][how-to tutorial], to visualize how the memory allocations/deallocations look like without having to recreate the UI.

The usage will be similar to PyTorch's:

• the user will be able to turn memory recording on and off (something like mx.record_memory_events(enabled=True|False))
• the user will be able to take a snapshot and dump it to a pickle file (mx.get_memory_snapshot() and mx.dump_memory_snapshot())
• The output pickle file generated from mx.dump_memory_snapshot() will be loadable directly in the web app

I've been studying how PyTorch implements it, and my current design (naming not finalized + Metal backend-only for now) is roughly:

• To turn memory recording on (and off):

  • a virtual function to turn on/off recording (including traceback capture) and configuration setup is added in the Allocator class (each backend-allocator, e.g., MetalAllocator will implement it)
  • a private bool member record_enabled is added to the MetalAllocator class; this is the flag to inform whether to create and store a new memory event

• To record memory events:

  • a private function maybe_record() is added to MetalAllocator, will be called in malloc/free-related functions to create a memory event and add it into an event buffer. The event creation + addition to the buffer only occur if record_enabled is true.
  • to obtain the context of each event (e.g., primitive name, traceback stack index) , a scope object is created right before gpu::eval(arr) / cpu::eval(arr) is run (mlx/transforms.cpp) that sets a thread_local current-op context. maybe_record() reads this context and adds it onto each memory event.
  • to capture a traceback (currently only Python stacks), array::ArrayDesc::init() (if capture is enabled) calls a hook that walks the current call stack (PyEval_GetFrame and PyFrame_GetBack), builds frames {code: PyObject*, line: int}, interns them, and returns an integer stack_idx to be stored on the ArrayDesc. During eval, the stack_idx will be copied into the thread_local current-op context I mentioned earlier

• To take a snapshot, a function will be called to symbolize the frames and add them into the memory events (and pass them to Python to be exported as a pickle file)

Writing this issue to check:

1. Would you be interested in this feature landing in MLX? (I'll open a PR once I'm done, if so.)
2. Would love thoughts on the design (e.g., things I should be aware of, recommendations on where it should live, particular naming conventions to follow).

Thank you so much in advance!

The main entry points mentioned are the Allocator and MetalAllocator classes, mlx/transforms.cpp, and array::ArrayDesc::init(). Start by tracing allocator malloc/free paths and the eval scope context, then compare the proposed APIs with PyTorch's memory-viz snapshot format. Done means recording can be toggled, snapshots can be exported, and the resulting pickle loads in the linked web app.