Skip to content

Commit 63a6f05

Browse files
QuafadasQuafadas
committed
ndarray plan
1 parent 8f2add5 commit 63a6f05

1 file changed

Lines changed: 334 additions & 0 deletions

File tree

‎site/notes/ndarray-design.md‎

Lines changed: 334 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,334 @@
1+
# NDArray Design Notes
2+
3+
## Critique of the Earlier Discussion
4+
5+
### What's right
6+
7+
1. **AD operates on whole tensors, not elements.** `Array[Jet[Double]]` is indeed catastrophic. `Jet[NDArray]` and `Node[NDArray]` are the correct shapes. This is the single most important insight.
8+
9+
2. **Runtime shapes over type-level dimension arithmetic.** Correct for a library targeting usability and cross-platform consistency. Type-level dimension encoding (`NDArray[3]`) leads to painful type-level arithmetic with reshape, broadcast, squeeze, and dynamic shapes. It doesn't compose well with Scala 3's type system without significant complexity.
10+
11+
3. **The interpreter pattern for AD.** Evaluating the same AST with different semantics (plain, forward AD, reverse AD) is clean and proven. Good separation of concerns.
12+
13+
4. **Contiguous `Array[Double]` as the memory substrate.** Correct for performance on all three platforms (JVM SIMD, JS typed arrays, Native CBLAS).
14+
15+
### What needs revision
16+
17+
#### 1. The `Tensor` type as proposed is too simplistic
18+
19+
```scala
20+
// From the discussion:
21+
case class Tensor(data: Array[Double], shape: Array[Int])
22+
```
23+
24+
This ignores **strides**, **offsets**, and **memory layout detection** — all things vecxt's `Matrix[A]` already handles. The existing Matrix has:
25+
26+
- `raw: Array[A]`, `rows`, `cols`, `rowStride`, `colStride`, `offset`
27+
- `isDenseColMajor`, `isDenseRowMajor`, `hasSimpleContiguousMemoryLayout`
28+
29+
An NDArray that can't represent views, transpositions, or slices without copying is a regression.
30+
31+
#### 2. "Tensor wraps NDArray" vs sharing a conceptual model
32+
33+
The discussion leans toward a single `Tensor` type that serves both as storage and as an AD computation node. These are fundamentally different concerns:
34+
35+
- **NDArray**: storage + shape + strides + element-wise and reduction ops. No gradient awareness.
36+
- **Tensor** (AD context): NDArray + gradient tracking + computation graph participation.
37+
38+
Conflating them means every NDArray carries AD overhead (even when you're just doing data manipulation), or you end up with two parallel APIs.
39+
40+
#### 3. `opaque type Vector = Tensor` / `opaque type Matrix = Tensor` is problematic
41+
42+
This loses:
43+
- `@specialized` (opaque types don't propagate specialization)
44+
- The rich existing Matrix API
45+
- Type-safe guarantees about dimensionality at construction time
46+
47+
Better: Matrix IS-A 2D NDArray (or wraps one with zero overhead), not an alias for an untyped Tensor.
48+
49+
#### 4. Element type genericity is missing
50+
51+
The discussion only mentions `Array[Double]`. Vecxt already supports `Matrix[Boolean]`, `Matrix[Int]`, and has `IntArrays`, `LongArrays`, `FloatArrays`. An NDArray should be `NDArray[A]` with specialization, not locked to Double.
52+
53+
#### 5. The MathAST / typeclass layer is a separate module concern
54+
55+
The AST, `Ops[A]` typeclass, and AD interpreters should NOT live in the NDArray module. NDArray is a data structure with operations. AD is a higher-level concern that *uses* NDArray. Mixing them creates circular dependencies and bloats the core.
56+
57+
---
58+
59+
## Revised Conceptual Model
60+
61+
### The shared abstraction: Strided N-dimensional Array
62+
63+
The core insight: **NDArray, Matrix, and (future) Tensor all share the same memory model:**
64+
65+
```
66+
┌─────────────────────────────────────────────┐
67+
│ Contiguous Array[A] in memory │
68+
│ ┌──────────────────────────────────────┐ │
69+
│ │ offset ──► shaped window into data │ │
70+
│ │ via strides │ │
71+
│ └──────────────────────────────────────┘ │
72+
└─────────────────────────────────────────────┘
73+
```
74+
75+
Every N-dimensional view is defined by:
76+
- `data: Array[A]` — the backing storage
77+
- `shape: Array[Int]` — dimensions `[d₀, d₁, ..., dₙ₋₁]`
78+
- `strides: Array[Int]` — memory step per dimension `[s₀, s₁, ..., sₙ₋₁]`
79+
- `offset: Int` — start position in `data`
80+
81+
Element at index `(i₀, i₁, ..., iₙ₋₁)` lives at:
82+
83+
$$\text{offset} + \sum_{k=0}^{n-1} i_k \cdot s_k$$
84+
85+
This is exactly what `Matrix[A]` already does for n=2:
86+
- `shape = [rows, cols]`
87+
- `strides = [rowStride, colStride]`
88+
89+
### Type hierarchy
90+
91+
```
92+
Array[A] ← 1D, no wrapper needed (vecxt's existing pattern)
93+
│
94+
NDArray[A] ← N-dimensional strided array (new)
95+
│
96+
Matrix[A] ← 2D NDArray with convenience API (migration from existing)
97+
```
98+
99+
Future (separate module):
100+
```
101+
NDArray[A]
102+
│
103+
Tensor ← NDArray[Double] + AD metadata (vecxt_ad module)
104+
```
105+
106+
### How Matrix becomes a 2D NDArray
107+
108+
**Not a breaking rewrite.** The plan:
109+
110+
1. `NDArray[A]` is introduced with `Array[A] + shape + strides + offset`
111+
2. `Matrix[A]` keeps its existing API but its internals become `ndim == 2` specialization of the same memory model
112+
3. Migration is gradual — Matrix can delegate to NDArray for shared ops (element access, reshape, transpose) while keeping its own ergonomic API (`@@` for matmul, `row()`, `col()`, etc.)
113+
114+
### How Tensor shares the model (future)
115+
116+
A future `Tensor` in a separate `vecxt_ad` module would be:
117+
118+
```scala
119+
case class Tensor(
120+
value: NDArray[Double], // the actual data
121+
// AD-specific fields:
122+
grad: Option[NDArray[Double]],
123+
backprop: NDArray[Double] => Unit
124+
)
125+
```
126+
127+
It doesn't wrap or duplicate NDArray — it *composes* with it. The NDArray is the storage; the Tensor adds gradient semantics. Same memory model, different concerns.
128+
129+
---
130+
131+
## NDArray[A] Detailed Design
132+
133+
### Core type
134+
135+
```scala
136+
class NDArray[@specialized(Double, Int, Float, Boolean) A] private[ndarray] (
137+
val data: Array[A],
138+
val shape: Array[Int],
139+
val strides: Array[Int],
140+
val offset: Int
141+
):
142+
lazy val ndim: Int = shape.length
143+
lazy val numel: Int = shape.product
144+
lazy val isContiguous: Boolean = /* check strides match dense layout */
145+
```
146+
147+
### Design decisions
148+
149+
| Decision | Choice | Rationale |
150+
|----------|--------|-----------|
151+
| Element type | Generic `A` with `@specialized` | Match existing Matrix pattern; avoid Double-only lock-in |
152+
| Shape representation | `Array[Int]` | Runtime, heap-efficient, no type-level arithmetic |
153+
| Strides | Explicit `Array[Int]` | Enables views, transpose, broadcast without copying |
154+
| Default layout | Column-major (F-order) for all ranks | Consistent with existing Matrix; BLAS-native; 2D slices of N-D arrays are directly BLAS-compatible; Julia validates this approach |
155+
| Constructor | Private + factory with BoundsCheck | Match existing Matrix pattern |
156+
| Bounds checking | Erasable `BoundsCheck` typeclass | Existing pattern, zero-cost in production |
157+
| API surface | Extension methods | Existing pattern, platform-specific dispatch |
158+
159+
### Layout convention
160+
161+
**Column-major (Fortran-order) for all ranks.** The first index varies fastest in memory.
162+
163+
For a dense array with `shape=[d₀, d₁, ..., dₙ₋₁]`, the default strides are `[1, d₀, d₀·d₁, ...]`.
164+
165+
Rationale:
166+
- Consistent with existing Matrix (column-major, BLAS-native)
167+
- 2D slices of N-D arrays are naturally column-major → directly BLAS-compatible, no layout conversion needed
168+
- Julia validates column-major-throughout as a performant, ergonomic choice for numerical computing
169+
- Row-major is still fully representable via different strides — nothing precludes it
170+
- Users who need row-major can construct with explicit strides or use a `rowMajor` factory
171+
172+
Detection methods (`isRowMajor`, `isColMajor`, `isContiguous`) enable fast-path optimizations regardless of layout.
173+
174+
### Relationship to existing code
175+
176+
| Existing | NDArray equivalent |
177+
|----------|-------------------|
178+
| `Matrix.raw` | `NDArray.data` |
179+
| `Matrix.rows` | `NDArray.shape(0)` |
180+
| `Matrix.cols` | `NDArray.shape(1)` |
181+
| `Matrix.rowStride` | `NDArray.strides(0)` |
182+
| `Matrix.colStride` | `NDArray.strides(1)` |
183+
| `Matrix.offset` | `NDArray.offset` |
184+
| `Matrix.hasSimpleContiguousMemoryLayout` | `NDArray.isContiguous` |
185+
| `Matrix.isDenseColMajor` | `NDArray.isColMajor` (for ndim==2) |
186+
| `Array[Double]` (1D ops) | Can be lifted to `NDArray` with `shape=[n], strides=[1]` |
187+
188+
### Broadcasting rules
189+
190+
Follow NumPy broadcasting semantics:
191+
1. Shapes are right-aligned
192+
2. Dimensions are compatible if equal or one of them is 1
193+
3. A dimension of 1 is broadcast (stride 0) to match the other
194+
195+
This is critical for AD (gradient accumulation involves broadcast reduction).
196+
197+
---
198+
199+
## Implementation Plan
200+
201+
### Milestone 0: Preparation (no code changes)
202+
- Agree on this design
203+
- alongside Matrix
204+
205+
### Milestone 1: NDArray core type + factories
206+
**Goal:** The type exists, can be constructed, and has basic properties.
207+
208+
- [ ] `NDArray[A]` class with `data`, `shape`, `strides`, `offset`
209+
- [ ] `@specialized(Double, Int, Float, Long, Boolean)`
210+
- [ ] Private constructor + `NDArray.apply(...)` with `BoundsCheck` validation
211+
- [ ] Stride validation (generalization of `strideMatInstantiateCheck`)
212+
- [ ] Factory methods: `NDArray.zeros`, `NDArray.ones`, `NDArray.fill`, `NDArray.fromArray` (1D), `NDArray.fromMatrix`
213+
- [ ] Properties: `ndim`, `numel`, `isContiguous`, `isRowMajor`, `isColMajor`
214+
- [ ] `toString` / `layout` for debugging
215+
- [ ] Cross-platform (shared `src/` only — no platform-specific code yet)
216+
- [ ] Unit tests for construction and property checking
217+
218+
**Deliverable:** `NDArray` can be created and inspected. No operations yet.
219+
220+
### Milestone 2: Indexing + views
221+
**Goal:** You can read/write elements and create views without copying.
222+
223+
- [ ] Single-element access: `ndarray(i, j, k)` → `A` (varargs `Int*` or `Array[Int]`)
224+
- [ ] Single-element update: `ndarray(i, j, k) = value`
225+
- [ ] Slice/view: `ndarray.slice(dim, range)` → new NDArray (adjusted strides/offset, no copy)
226+
- [ ] `transpose` (permute dimensions) → new NDArray (permuted strides, no copy)
227+
- [ ] `reshape` (contiguous arrays only, else error/copy)
228+
- [ ] `squeeze` (remove dimensions of size 1)
229+
- [ ] `unsqueeze` / `expandDims` (add dimension of size 1)
230+
- [ ] `flatten` → 1D NDArray (copy if non-contiguous)
231+
- [ ] `toArray` → contiguous `Array[A]` (copy if needed)
232+
- [ ] Cross-platform tests
233+
234+
**Deliverable:** Full indexing and view algebra. This is the foundation that everything else builds on.
235+
236+
### Milestone 3: Element-wise operations (Double)
237+
**Goal:** Arithmetic works for `NDArray[Double]`, with platform-specific acceleration.
238+
239+
- [ ] Binary ops: `+`, `-`, `*`, `/` (element-wise, with broadcasting)
240+
- [ ] Scalar ops: `ndarray + scalar`, `scalar * ndarray`, etc.
241+
- [ ] Unary ops: `neg`, `abs`, `exp`, `log`, `sqrt`, `tanh`, `sigmoid`
242+
- [ ] In-place variants: `+=`, `-=`, `*=`, `/=`
243+
- [ ] Comparison ops returning `NDArray[Boolean]`: `>`, `<`, `>=`, `<=`, `==`
244+
- [ ] Platform-specific implementations:
245+
- JVM: SIMD `DoubleVector` for contiguous arrays
246+
- JS/Native: while loops
247+
- [ ] Broadcasting implementation (shape compatibility check + stride-0 expansion)
248+
- [ ] Cross-platform tests with tolerance for floating point
249+
250+
**Deliverable:** NDArray is useful for numeric computation. Broadcasting works.
251+
252+
### Milestone 4: Reduction operations
253+
**Goal:** Aggregation along axes.
254+
255+
- [ ] Full reductions: `sum`, `mean`, `min`, `max`, `variance`
256+
- [ ] Axis reductions: `sum(axis=k)`, `mean(axis=k)`, etc. → NDArray with one fewer dimension
257+
- [ ] `argmin`, `argmax` (full and per-axis)
258+
- [ ] `dot` (1D), `matmul` (2D) — delegate to existing BLAS for 2D case
259+
- [ ] `norm` (L1, L2, Linf)
260+
- [ ] Platform-specific fast paths for contiguous data
261+
- [ ] Cross-platform tests
262+
263+
**Deliverable:** Statistical and aggregation workloads run on NDArray.
264+
265+
### Milestone 5: Matrix ↔ NDArray bridge
266+
**Goal:** Matrix and NDArray interoperate seamlessly.
267+
268+
- [ ] `Matrix.toNDArray` → 2D NDArray (zero-copy, shared backing array)
269+
- [ ] `NDArray.toMatrix` → Matrix (only if ndim==2, zero-copy)
270+
- [ ] `Array[Double].toNDArray` → 1D NDArray
271+
- [ ] Ensure existing Matrix operations still work (regression tests)
272+
- [ ] Consider: should Matrix internally delegate to NDArray for shared operations? (Evaluate performance implications first)
273+
- [ ] Documentation: migration guide for users
274+
275+
**Deliverable:** The two types coexist and convert freely. No breaking changes to Matrix API.
276+
277+
### Milestone 6: Extended element types + polish
278+
**Goal:** NDArray works for Int, Float, Boolean with appropriate ops.
279+
280+
- [ ] `NDArray[Int]`: arithmetic, comparisons, reductions
281+
- [ ] `NDArray[Boolean]`: logical ops (`&`, `|`, `!`), `any`, `all`, `countTrue`
282+
- [ ] `NDArray[Float]`: arithmetic (useful for ML workloads)
283+
- [ ] Boolean indexing: `ndarray(boolNdarray)` → filtered 1D result
284+
- [ ] `where(condition, x, y)` → element-wise conditional
285+
- [ ] Performance benchmarks vs existing `Array[Double]` ops
286+
- [ ] API review and cleanup
287+
288+
**Deliverable:** NDArray is a general-purpose N-dimensional array with full type support.
289+
290+
---
291+
292+
## Milestone dependency graph
293+
294+
```
295+
M0 (preparation)
296+
└─► M1 (core type)
297+
└─► M2 (indexing + views)
298+
├─► M3 (element-wise ops)
299+
│ └─► M4 (reductions)
300+
│ └─► M6 (extended types + polish)
301+
└─► M5 (Matrix bridge)
302+
```
303+
304+
M3 and M5 can proceed in parallel after M2.
305+
306+
---
307+
308+
## Future milestones (out of scope for NDArray, but planned)
309+
310+
### Milestone 7: AD module (`vecxt_ad`)
311+
- Separate Mill module depending on `vecxt`
312+
- `Tensor` type composing `NDArray[Double]` + gradient tracking
313+
- Forward-mode AD via `Jet[NDArray[Double]]`
314+
- Reverse-mode AD via `Node` with backprop closures
315+
- VJP rules for all NDArray primitives
316+
317+
### Milestone 8: MathAST + interpreters
318+
- Expression tree (`MathAST[A]`)
319+
- `Ops[A]` typeclass with instances for `NDArray`, `Jet[NDArray]`, `Node`
320+
- Plain evaluator, forward AD evaluator, reverse AD evaluator
321+
322+
---
323+
324+
## Open questions
325+
326+
1. ~~**Module placement**~~ → `vecxt/src/` (decided)
327+
328+
2. ~~**Default layout**~~ → Column-major (F-order) for all ranks (decided). Consistent with existing Matrix and BLAS.
329+
330+
3. **Copy semantics for views:** NDArray views share backing data. Mutation through one view is visible through others. This is the NumPy/PyTorch model and avoids unnecessary copies. Should we support copy-on-write? (Recommendation: no, too complex, just document the aliasing behavior.)
331+
332+
4. **Naming:** `NDArray[A]` vs `NdArray[A]` vs `Tensor[A]`? Recommendation: `NDArray` for the data structure. Reserve `Tensor` for the AD-aware type in a future module.
333+
334+
5. **Int indexing API:** Varargs `apply(indices: Int*)` is convenient but allocates. Alternative: overloads for 1, 2, 3, 4, N cases. Or: `IArray[Int]` to signal no mutation. Recommendation: specific overloads for 1-4D, varargs for N>4.

0 commit comments

Comments
 (0)