Add tensor package

3 files changed, 457 insertions, 0 deletions

diff --git a/tensor/dense.go b/tensor/dense.go
new file mode 100644
index 0000000..723cba5
--- /dev/null
+++ b/tensor/dense.go
@@ -0,0 +1,331 @@
+package tensor
+
+import (
+	"slices"
+)
+
+type Dense[T any] struct {
+	base    []T
+	offset  int
+	shape   []int
+	strides []int
+}
+
+func NewDense[T any](shape []int) Dense[T] {
+	r := len(shape)
+
+	if r == 0 {
+		return Dense[T]{
+			base:    make([]T, 1),
+			offset:  0,
+			shape:   make([]int, 0),
+			strides: make([]int, 0),
+		}
+	}
+
+	n := 1
+
+	for dim := range slices.Backward(shape) {
+		if shape[dim] < 0 {
+			panic("negative size")
+		}
+
+		n *= shape[dim]
+	}
+
+	strides := make([]int, len(shape))
+
+	for dim := range slices.Backward(strides) {
+		if dim == r-1 {
+			strides[dim] = 1
+
+			continue
+		}
+
+		strides[dim] = strides[dim+1] * shape[dim+1]
+	}
+
+	out := Dense[T]{
+		base:    make([]T, n),
+		offset:  0,
+		shape:   make([]int, len(shape)),
+		strides: strides,
+	}
+
+	copy(out.shape, shape)
+
+	return out
+}
+
+func (d Dense[T]) Size() int {
+	n := 1
+
+	for _, size := range d.shape {
+		n *= size
+	}
+
+	return n
+}
+
+func (d Dense[T]) Rank() int {
+	return len(d.shape)
+}
+
+func (d Dense[T]) Set(idxs []int, val T) {
+	storage := d.storageIndex(idxs)
+
+	d.base[storage] = val
+}
+
+func (d Dense[T]) At(idxs []int) T {
+	storage := d.storageIndex(idxs)
+
+	return d.base[storage]
+}
+
+func (d Dense[T]) IsEmpty() bool {
+	return d.Size() == 0
+}
+
+func (d Dense[T]) IsScalar() bool {
+	return d.Rank() == 0
+}
+
+func (d Dense[T]) IsContiguous() bool {
+	p := 1
+
+	for dim, stride := range slices.Backward(d.strides) {
+		if p != stride {
+			return false
+		}
+
+		p *= d.shape[dim]
+	}
+
+	return true
+}
+
+func (d Dense[T]) Contiguous() Dense[T] {
+	if d.IsContiguous() {
+		out := NewDense[T](d.shape)
+
+		copy(out.base, d.base[d.offset:d.offset+d.Size()])
+
+		return out
+	}
+
+	out := NewDense[T](d.shape)
+
+	idxs := make([]int, d.Rank())
+
+	linear := 0
+
+	for _, v := range d.All(idxs) {
+		out.base[linear] = v
+
+		linear++
+	}
+
+	return out
+}
+
+func (d Dense[T]) Permute(dims []int) Dense[T] {
+	r := d.Rank()
+
+	if len(dims) != r {
+		panic("invalid dimensions")
+	}
+
+	seen := make([]bool, r)
+
+	shape := make([]int, r)
+	strides := make([]int, r)
+
+	for out, in := range dims {
+		if in < 0 || in >= r {
+			panic("invalid dimension")
+		}
+
+		if seen[in] {
+			panic("duplicate dimension")
+		}
+
+		seen[in] = true
+
+		shape[out] = d.shape[in]
+		strides[out] = d.strides[in]
+	}
+
+	return Dense[T]{
+		base:    d.base,
+		offset:  d.offset,
+		shape:   shape,
+		strides: strides,
+	}
+}
+
+func (d Dense[T]) Slice(dim, start, end int) Dense[T] {
+	r := d.Rank()
+
+	if dim < 0 || dim >= r {
+		panic("invalid dimension")
+	}
+
+	if start < 0 || start > end || end > d.shape[dim] {
+		panic("invalid interval")
+	}
+
+	shape := make([]int, r)
+	strides := make([]int, r)
+
+	copy(shape, d.shape)
+	copy(strides, d.strides)
+
+	shape[dim] = end - start
+
+	return Dense[T]{
+		base:    d.base,
+		offset:  d.offset + start*strides[dim],
+		shape:   shape,
+		strides: strides,
+	}
+}
+
+func (d Dense[T]) Select(dim, index int) Dense[T] {
+	r := d.Rank()
+
+	if dim < 0 || dim >= r {
+		panic("invalid dimension")
+	}
+
+	if index < 0 || index >= d.shape[dim] {
+		panic("invalid index")
+	}
+
+	shape := make([]int, 0, len(d.shape)-1)
+	strides := make([]int, 0, len(d.strides)-1)
+
+	for dimension := range r {
+		if dimension == dim {
+			continue
+		}
+
+		shape = append(shape, d.shape[dimension])
+		strides = append(strides, d.strides[dimension])
+	}
+
+	return Dense[T]{
+		base:    d.base,
+		offset:  d.offset + index*d.strides[dim],
+		shape:   shape,
+		strides: strides,
+	}
+}
+
+func (d Dense[T]) All(idxs []int) func(func([]int, T) bool) {
+	r := d.Rank()
+	n := d.Size()
+
+	if len(idxs) != r {
+		panic("invalid indices")
+	}
+
+	return func(yield func([]int, T) bool) {
+		for i := range idxs {
+			idxs[i] = 0
+		}
+
+		storage := d.offset
+
+		for range n {
+			if !yield(idxs, d.base[storage]) {
+				return
+			}
+
+			for dim, size := range slices.Backward(d.shape) {
+				idxs[dim]++
+				storage += d.strides[dim]
+
+				if idxs[dim] < size {
+					break
+				}
+
+				idxs[dim] = 0
+				storage -= size * d.strides[dim]
+			}
+		}
+	}
+}
+
+func (d Dense[T]) unravel(linear int, idxs []int) {
+	r := d.Rank()
+	n := d.Size()
+
+	if len(idxs) != r {
+		panic("invalid indices")
+	}
+
+	if linear < 0 || linear >= n {
+		panic("invalid linear index")
+	}
+
+	for dim, size := range slices.Backward(d.shape) {
+		idxs[dim] = linear % size
+
+		linear /= size
+	}
+}
+
+func (d Dense[T]) ravel(idxs []int) int {
+	r := d.Rank()
+
+	if len(idxs) != r {
+		panic("invalid indices")
+	}
+
+	linear := 0
+
+	p := 1
+
+	for dim, size := range slices.Backward(d.shape) {
+		if idxs[dim] < 0 || idxs[dim] >= size {
+			panic("invalid index")
+		}
+
+		linear += idxs[dim] * p
+
+		p *= size
+	}
+
+	return linear
+}
+
+func (d Dense[T]) storageIndex(idxs []int) int {
+	r := d.Rank()
+
+	if len(idxs) != r {
+		panic("invalid indices")
+	}
+
+	storage := d.offset
+
+	for dim := range r {
+		if idxs[dim] < 0 || idxs[dim] >= d.shape[dim] {
+			panic("invalid index")
+		}
+
+		storage += idxs[dim] * d.strides[dim]
+	}
+
+	return storage
+}
+
+func (d Dense[T]) Data() ([]T, bool) {
+	if !d.IsContiguous() {
+		return nil, false
+	}
+
+	n := d.Size()
+
+	return d.base[d.offset : d.offset+n], true
+}
diff --git a/tensor/dense_test.go b/tensor/dense_test.go
new file mode 100644
index 0000000..5bc4e6f
--- /dev/null
+++ b/tensor/dense_test.go
@@ -0,0 +1,123 @@
+package tensor
+
+import (
+	"slices"
+	"testing"
+)
+
+func TestNewDense_Strides(t *testing.T) {
+	d := NewDense[float32]([]int{2, 3, 4})
+
+	expected := []int{12, 4, 1}
+
+	if !slices.Equal(d.strides, expected) {
+		t.Fatalf("expected %v but got %v", expected, d.strides)
+	}
+}
+
+func TestDense_IsContiguous(t *testing.T) {
+	d := NewDense[float32]([]int{2, 3})
+
+	if !d.IsContiguous() {
+		t.Fatalf("expected contiguous true")
+	}
+
+	permuted := d.Permute([]int{1, 0})
+
+	if permuted.IsContiguous() {
+		t.Fatalf("expected contiguous false")
+	}
+}
+
+func TestDense_IsContiguousStrict(t *testing.T) {
+	d := NewDense[int]([]int{1, 3})
+
+	v := Dense[int]{
+		base:    d.base,
+		offset:  0,
+		shape:   []int{1, 3},
+		strides: []int{999, 1},
+	}
+
+	if v.IsContiguous() {
+		t.Fatalf("expected contiguous false")
+	}
+}
+
+func TestDense_Contiguous(t *testing.T) {
+	d := NewDense[float32]([]int{2, 3})
+
+	permuted := d.Permute([]int{1, 0})
+
+	contiguous := permuted.Contiguous()
+
+	if !contiguous.IsContiguous() {
+		t.Fatalf("expected contiguous true")
+	}
+}
+
+func TestDense_Permute(t *testing.T) {
+	d := NewDense[float32]([]int{2, 3})
+
+	for i := range d.Size() {
+		d.base[i] = float32(i)
+	}
+
+	permuted := d.Permute([]int{1, 0})
+
+	expected := []float32{
+		0, 3,
+		1, 4,
+		2, 5,
+	}
+
+	buffer := make([]int, permuted.Rank())
+
+	linear := 0
+
+	for _, v := range permuted.All(buffer) {
+		a := v
+		b := expected[linear]
+
+		if a != b {
+			t.Fatalf("expected %.2f but got %.2f at linear %d", b, a, linear)
+		}
+
+		linear++
+	}
+}
+
+func TestDense_All(t *testing.T) {
+	d := NewDense[float32]([]int{2, 3})
+
+	for i := range d.Size() {
+		d.base[i] = float32(i)
+	}
+
+	d = d.Permute([]int{1, 0})
+
+	out := []float32{
+		0, 3,
+		1, 4,
+		2, 5,
+	}
+
+	buffer := make([]int, d.Rank())
+
+	step := 0
+
+	for idxs := range d.All(buffer) {
+		a := d.At(idxs)
+		b := out[step]
+
+		if a != b {
+			t.Fatalf("expected %.2f but got %.2f", a, b)
+		}
+
+		step++
+	}
+
+	if step != d.Size() {
+		t.Fatalf("premature termination")
+	}
+}
diff --git a/tensor/go.mod b/tensor/go.mod
new file mode 100644
index 0000000..765fa4f
--- /dev/null
+++ b/tensor/go.mod
@@ -0,0 +1,3 @@
+module tensor
+
+go 1.25.0