I've been working on neural network stuff in Nim and wanted to share this tensor code. Most people use Python, but Nim has some really good advantages for AI work.
Why Nim for AI?
# tensor.nim
# Educational, modular 2D tensor system for neural networks
import math, random, sequtils, strformat
# ------------------------
# Tensor Definition
# Core tensor structure with dimensions and 2D data storage
# ------------------------
type
Tensor* = object
rows*: int
cols*: int
data*: seq[seq[float]]
# ------------------------
# Tensor Constructors
# Factory functions for creating tensors with different initialization patterns
# ------------------------
proc zeros*(rows, cols: int): Tensor =
result.rows = rows
result.cols = cols
result.data = newSeqWith(rows, newSeqWith(cols, 0.0))
proc ones*(rows, cols: int): Tensor =
result.rows = rows
result.cols = cols
result.data = newSeqWith(rows, newSeqWith(cols, 1.0))
proc randomNormal*(rows, cols: int): Tensor =
## Xavier initialization (mean=0, std ~ sqrt(2 / (in + out)))
## Uses Gaussian distribution for better neural network weight initialization
result.rows = rows
result.cols = cols
result.data = newSeq[seq[float]](rows)
let scale = sqrt(2.0 / float(rows + cols))
for i in 0..<rows:
result.data[i] = newSeq[float](cols)
for j in 0..<cols:
result.data[i][j] = gauss(0.0, scale)
proc makeTensor*(data: seq[seq[float]]): Tensor =
result.rows = data.len
result.cols = if data.len > 0: data[0].len else: 0
result.data = data
# ------------------------
# Shape & Access
# Utility functions for getting tensor dimensions and accessing elements
# ------------------------
proc shape*(t: Tensor): (int, int) =
(t.rows, t.cols)
proc `[]`*(t: Tensor, row, col: int): float =
t.data[row][col]
proc `[]=`*(t: var Tensor, row, col: int, value: float) =
t.data[row][col] = value
# ------------------------
# Basic Element-wise Math
# Fundamental arithmetic operations for tensors and scalars
# ------------------------
proc `+`*(a, b: Tensor): Tensor =
assert a.rows == b.rows and a.cols == b.cols
result = zeros(a.rows, a.cols)
for i in 0..<a.rows:
for j in 0..<a.cols:
result[i, j] = a[i, j] + b[i, j]
proc `+`*(t: Tensor, scalar: float): Tensor =
result = zeros(t.rows, t.cols)
for i in 0..<t.rows:
for j in 0..<t.cols:
result[i, j] = t[i, j] + scalar
proc `-`*(a, b: Tensor): Tensor =
assert a.rows == b.rows and a.cols == b.cols
result = zeros(a.rows, a.cols)
for i in 0..<a.rows:
for j in 0..<a.cols:
result[i, j] = a[i, j] - b[i, j]
proc `*`*(a, b: Tensor): Tensor =
## Hadamard product (element-wise multiplication)
## This is NOT matrix multiplication - use dot() for that
assert a.rows == b.rows and a.cols == b.cols
result = zeros(a.rows, a.cols)
for i in 0..<a.rows:
for j in 0..<a.cols:
result[i, j] = a[i, j] * b[i, j]
proc `*`*(t: Tensor, scalar: float): Tensor =
result = zeros(t.rows, t.cols)
for i in 0..<t.rows:
for j in 0..<t.cols:
result[i, j] = t[i, j] * scalar
# ------------------------
# Matrix Operations
# Linear algebra operations including matrix multiplication and transpose
# ------------------------
proc dot*(a, b: Tensor): Tensor =
assert a.cols == b.rows
result = zeros(a.rows, b.cols)
for i in 0..<a.rows:
for j in 0..<b.cols:
var sum = 0.0
for k in 0..<a.cols:
sum += a[i, k] * b[k, j]
result[i, j] = sum
proc transpose*(t: Tensor): Tensor =
result = zeros(t.cols, t.rows)
for i in 0..<t.rows:
for j in 0..<t.cols:
result[j, i] = t[i, j]
proc apply*(t: Tensor, fn: proc(x: float): float): Tensor =
## Apply a function element-wise to the tensor
## Useful for activation functions like ReLU, sigmoid, etc.
result = zeros(t.rows, t.cols)
for i in 0..<t.rows:
for j in 0..<t.cols:
result[i, j] = fn(t[i, j])
# ------------------------
# Tensor Utilities
# Statistical and utility functions for tensor manipulation
# ------------------------
proc sum*(t: Tensor): float =
result = 0.0
for i in 0..<t.rows:
for j in 0..<t.cols:
result += t[i, j]
proc mean*(t: Tensor): float =
t.sum() / float(t.rows * t.cols)
proc copy*(t: Tensor): Tensor =
## Create a deep copy of the tensor
result = zeros(t.rows, t.cols)
for i in 0..<t.rows:
for j in 0..<t.cols:
result[i, j] = t[i, j]
proc `+=`*(a: var Tensor, b: Tensor) =
## In-place addition - modifies the first tensor
assert a.rows == b.rows and a.cols == b.cols
for i in 0..<a.rows:
for j in 0..<a.cols:
a[i, j] = a[i, j] + b[i, j]
proc `-=`*(a: var Tensor, b: Tensor) =
## In-place subtraction - modifies the first tensor
assert a.rows == b.rows and a.cols == b.cols
for i in 0..<a.rows:
for j in 0..<a.cols:
a[i, j] = a[i, j] - b[i, j]
proc `$`*(t: Tensor): string =
## String representation for pretty printing tensors
result = &"Tensor({t.rows}x{t.cols}):\n"
for i in 0..<t.rows:
result.add " ["
for j in 0..<t.cols:
result.add &"{t[i, j]:>7.4f}"
if j < t.cols - 1: result.add ", "
result.add "]\n"
# ------------------------
# Simple Test Block
# Demonstrates basic tensor operations and functionality
# ------------------------
when isMainModule:
echo "--- Tensor Test ---"
let a = randomNormal(2, 3)
let b = ones(2, 3)
let c = a + b
let d = a * 2.0
let e = dot(a, transpose(b))
echo "Tensor A:"
echo a
echo "Tensor B:"
echo b
echo "A + B:"
echo c
echo "A * 2.0:"
echo d
echo "A dot B^T:"
echo e
proc relu(x: float): float =
if x > 0: x else: 0.0
echo "ReLU(A):"
echo apply(a, relu)
var f = copy(a)
f += b
echo "Copy and += test:"
echo f
var g = copy(b)
g -= ones(2, 3) * 0.3
echo "Copy and -= test:"
echo gThis does all the math stuff you need for neural networks. The Tensor thing holds your data in rows and columns. You can make tensors full of zeros, ones, or random numbers that work good for training.
The math works like you'd expect - a + b adds them together. The * does element by element multiplication while dot() does real matrix multiplication. The apply() function lets you run things like ReLU on every number in your tensor.
It has all the helper stuff too - adding up all numbers, getting averages, copying tensors, and printing them out nice. The test part shows it actually working with real math you'd use.
This is what I'm building from. Next I'm adding layers and training stuff. Goal is to solve XOR and maybe MNIST later, all with code that's easy to read and compiles to one fast file.
Anyone else trying Nim for machine learning? It feels like it should be more popular for this stuff.
look into arraymancer instead of defining your own structure
No harm in building your own, especially when learning.
exprgrad is also a nice existing library