How I Trained a Neural Network in Nushell
tl;dr
I trained a neural network in Nushell by wrapping libtorch, the same C++ library that powers PyTorch, in a Nushell plugin, making it possible to create and process tensors on the command line.
Furthermore, I wrapped an entire web app and web browser in a second Nushell plugin to make rendering beautiful, high resolution plots in the terminal possible.
Together, I used these tools to train a neural network to classify data and visualize the results, all inside my terminal.
I also compared the results to PyTorch and Matplotlib, which match exactly.
Results Preview
The following plot illustrates the neural network that I created with Nushell. I started by generating data points that fall into three buckets. I then used a two-layer neural network to classify the points. The final plot shows both the original “true” data points, and a contour map of the model, showing the model’s predictions, and how well they match the original data.
What is Nushell?
Nushell is a new shell that is designed to handle structured data, unlike traditional shells like zsh and bash, which are focused on text. It is written in Rust and has many features of modern programming languages, including a powerful type system, good errors, and a plugin system.
Why Nushell?
Although Python is a good programming language, I always thought it would be cool if I could pipe data into my GPU directly from the command line. Because Python is not a shell, this is impractical with Python. However, Nushell is a shell, so it is designed to handle structured data and can easily pipe data from one command to another.
Although building a similar proof of concept with zsh or bash is possible, no one would ever want to do data analysis in those languages, because they are too slow for all import/export operations. Nushell, however, handles structured data efficiently, so it is sufficiently fast to pipe complex data from command to command. The tooling for Nushell may not yet be as mature as Python, but in principle it is just as capable of data analysis as Python, and brings a whole new dimension of value to the table by being your shell, and thus working out of the box with all of your existing shell commands and workflows.
How Does it Work?
I created two tools to make this possible: Nutorch and Termplot.
Nutorch
Nutorch is a Nushell plugin written in Rust that wraps tch-rs, which itself is a Rust wrapper for libtorch, the C++ library that powers PyTorch. This plugin allows you to create and manipulate tensors in Nushell, which are the fundamental data structure used in neural networks (and a lot of other data analysis). The API is designed to be similar to PyTorch, so if you are familiar with PyTorch, you will feel right at home. It also has some convenient Nushell-specific features, such as the ability to pipe data directly into tensors from other commands.
The fundamental idea of Nutorch is very simple: You can create tensors, either by piping data into it or generating them from scratch, on any device (CPU or CUDA or MPS). Those tensors are stored in a hash table which lives as long as the Nushell plugin needs to live. Nushell decides when the plugin should exit, which by default lasts 10 seconds since the last command.
Commands then accept tensors as input or arguments, and the tensors are
retrieved from memory and operated on. To view the final result, you can run
torch value which will download the tensor and convert it into a Nushell
value, which can then be printed or piped to another command.
Termplot
Termplot is a Node.js CLI tool that runs an entire headless web browser (via puppeteer) with a web server running a React Router web app with an in-memory database of JSON configuration files to render plotly.js plots (or, in the future, any other web-based plotting library). Termplot takes a screen shot of the plot and renders it to the terminal using ANSI escape codes.
The standard Termplot CLI tool must open an entire web browser for every plot, which is slow. So I also made a Nushell plugin for Termplot that manages the web browser in the background, making plots render almost instantly. The normal CLI tool takes about 2 seconds to render a plot, but the Nushell plugin renders plots in about 20 milliseconds.
beautiful.nu
I have also created a simple Nushell script called beautiful.nu that generates catppuccin-themed plotly.js configuration files, thus bridging Nutorch and Termplot. The beautiful.nu code can be found on NPM here and on GitHub in the Nutorch repo here.
Simple One-Liner Example
As a simple example of the power of Notorch and Termplot, considering the following powerful one-liner:
torch arange 0 10 0.1 --device mps | torch exp | torch value | [{y: $in}] | beautiful lines | termplot
This command generates a tensor of values from 0 to 10 in increments of 0.1,
computes the exponential of each value, downloads the result from the GPU,
generates a “beautiful” plot configuration from the data, and then plots the
result using Termplot. The --device mps flag specifies that the computation
should be done on the Apple M1 GPU (if available). The final result is a
beautiful plot of the exponential function rendered directly in the terminal.
Have you ever seen a one-liner like this does does a computation on the GPU, and renders the result directly in the terminal? This is the power of Nutorch and Termplot working together.
Neural Network Example
I trained a neural network in Nushell using Nutorch and Termplot in the following manner:
- I generated a series of data points that fall into three separate clusters. Each data point corresponds to a label, which is the cluster it belongs to. The labels are integers 0, 1, or 2.
- Knowing it was possible to use a two-layer neural network with ReLU activation to classify the data, I created a Nushell function that initializes the model parameters (weights and biases) for a two-layer neural network with ReLU activation.
- I created a Nushell function that performs a forward pass through the model, which computes the logits (raw outputs) of the model.
- I created a Nushell function that computes the cross entropy loss between the logits and the target labels.
- I created a Nutorch function, called from Nushell, that performs stochastic gradient descent (SGD) to update the model parameters based on the loss.
- I created a Nushell function that trains the model for a specified number of epochs, recording the loss every so often.
- I created a Nushell function that plots the raw data points using Termplot.
- I created a Nushell function that plots the loss over time using Termplot.
- I created a Nushell function that plots the model predictions using Termplot, which shows the decision boundary of the model and how well it matches the original data.
- I programmed the exact same logic in Python using PyTorch and Matplotlib to verify that the results match exactly. Because the underlying library is libtorch, using the same random seed, the results are exactly the same, proving the Nushell code is working exactly as intended. The Python code was written first and used as a template to make sure the Nushell code was complete.
You can see the code for the Nushell neural network example and the corresponding code in Python in the appendix below. You can also find additional information on installing Nutorch and Termplot on their respective GitHub pages:
Visualizing the Results
The following plot shows the initial data points, colored by their class.
The following plot shows the loss over time during training.
The following plot shows the model predictions, with the decision boundary overlaid on the original data points.
These plots were created using Nutorch and Termplot, but similar plots can be produced using the Python code in the appendix below.
Conclusion
Modern technologies including Nushell, libtorch, and tch-rs make it possible to do data analysis and machine learning in a shell environment. This is extremely convenient and powerful for anyone who spends a lot of time in the terminal because it is compatible with all of your existing shell commands and workflows.
A proof of concept neural network was created in Nushell using Nutorch and Termplot, which demonstrates the power of these tools. The neural network was trained to classify data points into three clusters, and the results were visualized using Termplot. The results match exactly with a similar implementation in Python using PyTorch and Matplotlib, proving that the Nushell implementation is working as intended.
To learn more about how to use Nutorch and Termplot, please visit their respective websites, which currently redirect to GitHub:
Appendix A: Nushell Neural Network Example
plugin use torch
plugin use termplot
source node_modules/termplot.nu/termplot.nu
use node_modules/beautiful.nu *
torch manual_seed 42
def generate_data [
--n_samples: int = 300 # Number of samples to generate
--centers: int = 3 # Number of cluster centers
--cluster_std: float = 0.7 # Standard deviation of clusters
--skew_factor: float = 0.3 # Skew factor for data distribution
]: [nothing -> record<X: string, y: string>] {
let n_samples_per_class: int = ($n_samples // $centers)
mut X_list: list<string> = [] # nutorch tensors have string ids
mut y_list: list<string> = [] # nutorch tensors have string ids
let blob_centers: list<string> = [
(torch tensor [0.0 0.0])
(torch tensor [3.0 0.0])
(torch tensor [1.5 2.5])
]
for i in (seq 0 ($centers - 1)) {
mut points: string = (torch randn $n_samples_per_class 2) | torch mul (torch tensor $cluster_std) | torch add ($blob_centers | get $i)
if $i == 1 or $i == 2 {
let center = ($blob_centers | get $i)
let skew = torch tensor [[1.0 ($skew_factor * ($i - 1))] [($skew_factor * ($i - 1)) 1.0]]
$points = $points | torch sub $center | torch mm $skew | torch add $center
}
let labels: string = torch full [$n_samples_per_class] $i --dtype 'int64'
$X_list = $X_list | append $points
$y_list = $y_list | append $labels
}
let X: string = $X_list | torch cat --dim 0
let y: string = $y_list | torch cat --dim 0
{X: $X y: $y}
}
def model_init [
--input_size: int = 2 # Number of input features
--hidden_size: int = 20 # Number of hidden units
--output_size: int = 3 # Number of output classes
]: [nothing -> record<w1: string, b1: string, w2: string, b2: string>] {
{
w1: (torch randn $hidden_size $input_size --requires_grad true)
b1: (torch randn $hidden_size --requires_grad true)
w2: (torch randn $output_size $hidden_size --requires_grad true)
b2: (torch randn $output_size --requires_grad true)
}
}
def model_get_parameters [
--model: record<w1: string, b1: string, w2: string, b2: string>
]: [nothing -> list<string>] {
[$model.w1 $model.b1 $model.w2 $model.b2]
}
def model_forward_pass [
--model: record<w1: string, b1: string, w2: string, b2: string>
]: [string -> string] {
torch mm ($model.w1 | torch t) # Matrix multiplication with input and first layer weights
| torch add $model.b1 # Add bias for first layer
| torch maximum ([0.0] | torch tensor) # ReLU activation
| torch mm ($model.w2 | torch t) # Matrix multiplication with second layer weights
| torch add $model.b2 # Add bias for second layer
}
def cross_entropy_loss [
--logits: string # tensor id of model outputs
--targets: string # tensor id of target labels
]: [nothing -> string] {
let logp = $logits | torch log_softmax --dim 1
# print $"logp: ($logp | torch mean | torch value)"
let loss = $logp | torch gather 1 ($targets | torch unsqueeze 1) | torch squeeze 1 | torch mean | torch neg
$loss
}
def train [
--model: record<w1: string, b1: string, w2: string, b2: string>
--X: string # Input tensor id
--y: string # Target tensor id
--epochs: int = 1000
--lr: float = 0.1
--record_every: int = 100
]: [nothing -> record<model: record<w1: string, b1: string, w2: string, b2: string>, losses: list<number>, steps: list<number>>] {
mut losses: list<number> = []
mut steps: list<number> = []
let ps = model_get_parameters --model $model
for epoch in (seq 0 ($epochs - 1)) {
let logits = $X | model_forward_pass --model $model
let loss = cross_entropy_loss --logits $logits --targets $y
for p in $ps {
$p | torch zero_grad
}
$loss | torch backward
torch sgd_step $ps --lr $lr
if ($epoch + 1) mod $record_every == 0 {
$losses = $losses | append ($loss | torch value)
$steps = $steps | append ($epoch + 1)
print $"epoch: ($epoch + 1)/($epochs), loss: (($loss | torch value | math round --precision 4))"
}
}
return {
model: $model
losses: $losses
steps: $steps
}
}
def plot_raw_data [res: record<X: string, y: string>] {
# Call with named arguments (flags)
let X: string = $res.X
let y: string = $res.y
let X_value = $X | torch value
let y_value = $y | torch value
[
{
x: ($X_value | enumerate | each {|xy| if ($y_value | get $xy.index) == 0 { $xy.item.0 } })
y: ($X_value | enumerate | each {|xy| if ($y_value | get $xy.index) == 0 { $xy.item.1 } })
}
{
x: ($X_value | enumerate | each {|xy| if ($y_value | get $xy.index) == 1 { $xy.item.0 } })
y: ($X_value | enumerate | each {|xy| if ($y_value | get $xy.index) == 1 { $xy.item.1 } })
}
{
x: ($X_value | enumerate | each {|xy| if ($y_value | get $xy.index) == 2 { $xy.item.0 } })
y: ($X_value | enumerate | each {|xy| if ($y_value | get $xy.index) == 2 { $xy.item.1 } })
}
] | beautiful scatter | merge deep {layout: {title: {text: "Raw Data"}}} | termplot
}
def plot_loss [
--losses: list<number> # list of loss values
--steps: list<number> # list of steps (epochs) corresponding to losses
] {
[{x: $steps y: $losses}] | beautiful lines | merge deep {layout: {title: {text: "Loss"}}} | termplot
}
def plot_results [
--X: string # Input tensor id
--y: string # Target tensor id
--model: record<w1: string, b1: string, w2: string, b2: string> # Model parameters
]: [nothing -> nothing] {
let Xl = $X | torch detach | torch value
let yl = $y | torch detach | torch value
# let yscaledl = $y | torch div ($y | torch max) | torch detach | torch value
let x_min = ($Xl | each {|x| $x | get 0 }) | math min
let x_max = ($Xl | each {|x| $x | get 0 }) | math max
let y_min = ($Xl | each {|x| $x | get 1 }) | math min
let y_max = ($Xl | each {|x| $x | get 1 }) | math max
let xs = torch arange $x_min $x_max 0.1
# let ys = torch arange $y_min $y_max 0.1
let ys = $xs
let mesh = torch stack [
($xs | torch repeat ($ys | torch value | length))
($ys | torch repeat_interleave ($xs | torch value | length))
] --dim 1
let logits = $mesh | model_forward_pass --model $model
let Z = torch argmax $logits --dim 1 | torch reshape [($xs | torch value | length) ($ys | torch value | length)]
beautiful plot
| beautiful add contour {
x: ($xs | torch value)
y: ($ys | torch value)
z: ($Z | torch value)
colorscale: (beautiful colorscale 3)
opacity: 0.4
}
| beautiful add scatter {
x: ($Xl | enumerate | each {|xy| if (($yl | get $xy.index) == 0) { $xy.item.0 } })
y: ($Xl | enumerate | each {|xy| if (($yl | get $xy.index) == 0) { $xy.item.1 } })
}
| beautiful add scatter {
x: ($Xl | enumerate | each {|xy| if (($yl | get $xy.index) == 1) { $xy.item.0 } })
y: ($Xl | enumerate | each {|xy| if (($yl | get $xy.index) == 1) { $xy.item.1 } })
}
| beautiful add scatter {
x: ($Xl | enumerate | each {|xy| if (($yl | get $xy.index) == 2) { $xy.item.0 } })
y: ($Xl | enumerate | each {|xy| if (($yl | get $xy.index) == 2) { $xy.item.1 } })
}
| merge deep {layout: {title: {text: "Model Predictions"}}} | termplot
}
let raw_data = generate_data --n_samples 300 --centers 3 --cluster_std 0.7 --skew_factor 0.3
plot_raw_data $raw_data
let net = model_init --input_size 2 --hidden_size 20 --output_size 3
let model_res = train --model $net --X $raw_data.X --y $raw_data.y --epochs 3000 --lr 0.1 --record_every 100
plot_loss --losses $model_res.losses --steps $model_res.steps
plot_results --X $raw_data.X --y $raw_data.y --model $model_res.model
Appendix B: Python Neural Network Example
import torch
from typing import List, Tuple, Dict
import matplotlib.pyplot as plt
torch.manual_seed(42) # reproducibility
def generate_data(
n_samples: int = 300,
centers: int = 3,
cluster_std: float = 0.7,
skew_factor: float = 0.3,
) -> Tuple[torch.Tensor, torch.Tensor]:
n_per_class = n_samples // centers
X_parts, y_parts = [], []
blob_centers = [
torch.tensor([0.0, 0.0]),
torch.tensor([3.0, 0.0]),
torch.tensor([1.5, 2.5]),
]
for i in range(centers):
pts = torch.randn(n_per_class, 2) * cluster_std + blob_centers[i]
if i in (1, 2):
skew = torch.tensor(
[[1.0, skew_factor * (i - 1)], [skew_factor * (i - 1), 1.0]]
)
pts = torch.mm(pts - blob_centers[i], skew) + blob_centers[i]
X_parts.append(pts)
y_parts.append(torch.full((n_per_class,), i, dtype=torch.long))
return torch.cat(X_parts), torch.cat(y_parts)
Model = Dict[str, torch.Tensor]
def model_init(inp: int = 2, hid: int = 20, out: int = 3) -> Model:
return {
"w1": torch.randn(hid, inp, requires_grad=True),
"b1": torch.randn(hid, requires_grad=True),
"w2": torch.randn(out, hid, requires_grad=True),
"b2": torch.randn(out, requires_grad=True),
}
def model_get_parameters(model: Model) -> List[torch.Tensor]:
return [model["w1"], model["b1"], model["w2"], model["b2"]]
def model_forward_pass(model: Model, x: torch.Tensor) -> torch.Tensor:
w1t = model["w1"].t()
x = torch.mm(x, w1t) + model["b1"]
x = torch.max(torch.tensor(0.0), x) # ReLU
w2t = model["w2"].t()
x = torch.mm(x, w2t) + model["b2"]
return x
def cross_entropy_loss(logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
logp = torch.log_softmax(logits, dim=1)
# print(f"logp: {logp.mean()}, targets: {targets.shape}")
chosen = torch.gather(logp, 1, targets.unsqueeze(1)).squeeze(1)
return -chosen.mean()
def sgd_step(ps: List[torch.Tensor], lr: float = 0.1) -> None:
"""
Vanilla gradient descent: p ← p - lr * p.grad , then reset gradients.
Operates in-place; returns nothing.
"""
with torch.no_grad():
for p in ps:
if p.grad is not None:
p -= lr * p.grad
def train(
model: Model,
X: torch.Tensor,
y: torch.Tensor,
epochs: int = 1000,
lr: float = 0.1,
record_every: int = 100,
) -> Tuple[List[float], List[int]]:
losses, steps = [], []
ps = model_get_parameters(model)
for epoch in range(epochs):
# forward & loss
logits = model_forward_pass(model, X)
loss = cross_entropy_loss(logits, y)
# zero existing grads, back-prop, SGD update
for p in ps:
if p.grad is not None:
p.grad.zero_()
loss.backward()
sgd_step(ps, lr)
if (epoch + 1) % record_every == 0:
losses.append(loss.item())
steps.append(epoch + 1)
print(f"epoch {epoch+1:4d}/{epochs} loss {loss.item():.4f}")
return losses, steps
def plot_raw_data(X: torch.Tensor, y: torch.Tensor) -> None:
Xl, yl = X.tolist(), y.tolist()
plt.scatter([p[0] for p in Xl], [p[1] for p in Xl], c=yl, alpha=0.8, cmap="viridis")
plt.title("Raw data")
plt.show()
def plot_loss(losses: List[float], steps: List[int]) -> None:
plt.plot(steps, losses)
plt.title("Training loss")
plt.xlabel("epoch")
plt.ylabel("loss")
plt.show()
def plot_results(X: torch.Tensor, y: torch.Tensor, model: Model) -> None:
Xl = X.detach().tolist()
yl = y.detach().tolist()
x_min = min(p[0] for p in Xl) - 1
x_max = max(p[0] for p in Xl) + 1
y_min = min(p[1] for p in Xl) - 1
y_max = max(p[1] for p in Xl) + 1
xs = torch.arange(x_min, x_max, 0.1)
ys = torch.arange(y_min, y_max, 0.1)
mesh = torch.stack([xs.repeat(len(ys)), ys.repeat_interleave(len(xs))], dim=1)
# note: do not use no_grad here for easier translating to nushell
logits = model_forward_pass(model, mesh)
Z = torch.argmax(logits, dim=1).reshape(len(ys), len(xs))
plt.contourf(xs, ys, Z, alpha=0.4, cmap="viridis")
plt.scatter([p[0] for p in Xl], [p[1] for p in Xl], c=yl, alpha=0.8, cmap="viridis")
plt.title("Decision boundary")
plt.show()
if __name__ == "__main__":
X, y = generate_data(n_samples=300, centers=3, cluster_std=0.7, skew_factor=0.3)
plot_raw_data(X, y)
net = model_init(inp=2, hid=20, out=3)
losses, steps = train(net, X, y, epochs=3000, lr=0.1, record_every=100)
plot_loss(losses, steps)
plot_results(X, y, net)