dispy.nvim is a fun little Neovim plugin I wrote to display PyTorch tensors directly in Neovim via the kitty graphics protocol. This writeup is intended to provide a brief overview of my thought process behind the implementation and motivate the plugin’s creation.
Table of contents
Open Table of contents
Motivation
I’m currently a PhD student at Virginia Tech studying machine learning, so I frequently work with PyTorch tensors. While PyTorch is great, working with high-dimensional data can get confusing and I often find myself wanting the ability to quickly sanity check specific tensors. For example, if I’m working with an image dataset, I want to be able to quickly sample and plot example images from the dataset. Additionally, when I’m training a model, I want to be able to calculate statistics from the model outputs and intermediate model activations without breaking flow.
However, I couldn’t find a great way to do this. I do most of my development work in Neovim using nvim-dap and debugpy to debug Python code. One of the nice features nvim-dap provides is a built-in REPL, which lets me execute code while debugging. It’s really helpful when I want to visualize the program state or make causal interventions. While the REPL is incredibly handy, I quickly got tired of manually importing matplotlib, saving off tensors as images, and manually viewing them using the icat kitten every time I wanted to analyze the state of my program (I’ve found the interactive matplotlib backends to be somewhat inconsistent and finicky in Neovim). nvim-dap has a built-in hover() method, but the output it produces when targeting a tensor is text-only and very cluttered, making it tricky to find what I’m looking for at a glance.
:lua require('dap').hover() - way too much information! My Solution
One of the aspects about Neovim I appreciate most is its configurability. Because getting started writing custom plugins in Lua is straightforward, Neovim has a vibrant ecosystem of user-created plugins. I’ve spent considerable time experimenting with various external plugins tailoring my editor to meet my needs (including adapting existing plugins), but this was my first experience building a project from scratch in Lua.
To visualize a user workflow, I put together a quick flow graph using Excalidraw. I wanted the user to be able to be able to choose a tensor, call a dispy.nvim method (preferably via keyboard shortcuts), and then see the requested data in the specified form (potentially a single image, multiple random images, a histogram of tensor values, or something else).
dispy.nvim user. Implementation Details
How do I go about actually building a plugin to do this? As mentioned above, I use nvim-dap as my Debug Adapter Protocol client of choice - this is how I debug Python programs, so anything I create has to be able to talk to nvim-dap somehow. Luckily, nvim-dap gives us programmatic access to a REPL we can use to query the currently running program state. Once the user pauses program execution (by hitting a breakpoint), dispy.nvim can talk to the REPL to execute Python code within the context of the paused program. This means we can indirectly use Python libraries such as matplotlib to plot the tensor as an image or lovely-tensors to quickly display information about the tensor.
Won't executing code by calling the REPL mess up my program state?
While arbitrary code execution via the REPL can definitely lead to unintentional, hard-to-reproduce bugs and / or incorrect output, I tried to write dispy.nvim in such a way that it affects the state of the currently executing program as little as possible. Also, dispy.nvim doesn’t send anything to the REPL except when directly triggered by the user. I think it’s a worthwhile tradeoff in my own work, but if there’s a better way to interact with the running program that gets around these concerns, please leave an issue on the GitHub repository and let me know!
To display the generated plots, dispy.nvim makes use of image.nvim under the hood. I recently started using image.nvim, which integrates with the Kitty graphics protocol, because I wanted to be able to see content images while writing markdown documents. Extending that use case was a primary motivator for dispy.nvim.
The demo video below gives a quick overview of some of the current functionality of the plugin.
Future work
There are still a number of features I’d like to see in dispy.nvim in the future, such as extending functionality directly to visualize model weights and reworking a few “hacky” solutions that went into getting the plugin working. I also still need to test on other platforms as well (only tested on MacOS so far) and terminals (Ghostty and Kitty both work, as well as theoretically other applications supporting the Kitty graphics protocol).
Conclusion
I often have ideas that I think are interesting and worth pursuing. I sit down at my computer, make a little bit of progress, get sidetracked trying to make my code perfect, and utterly fail to produce anything meaningful. Recently, I read an excellent article that beautifully expressed a similar sentiment. I know my code isn’t perfect. Neither is my writing, as you’ve certainly noticed by now. However, by being consistent about open-sourcing code and writing long-form content, I’m hoping to, as the author puts it, “slowly improve until I reach a point where I can be proud of my work.”
If you liked this writeup, please head over to GitHub repo for dispy.nvim and give it a star!