HEAL memory management diagram
Example
```plaintext A.shape = [ 3, 1, 5] B.shape = [2, 1, 3, 1, 1] ``` Becomes (after left-padding `A` with two 1s): ```plaintext A.shape = [1, 1, 3, 1, 5] B.shape = [2, 1, 3, 1, 1] ``` Resulting broadcasted shape: ```plaintext [2, 1, 3, 1, 5] ```Why is this a valid example?
To align shapes, we left-pad B’s shape: ``` Shape A: [5, 1, 4, 1] Shape B: [1, 3, 1, 1] ``` Comparison (from last to first): | Dimension | A | B | Valid? | | --------- | - | - | ---------- | | 4th | 1 | 1 | ✓ equal | | 3rd | 4 | 1 | ✓ one is 1 | | 2nd | 1 | 3 | ✓ one is 1 | | 1st | 5 | 1 | ✓ one is 1 | ✓ All dimensions are compatible → **broadcasting is allowed**Why is this an invalid example?
Left-padding B’s shape: ``` Shape A: [5, 2, 4, 1] Shape B: [1, 3, 1, 1] ``` Comparison: | Dimension | A | B | Valid? | | --------- | - | - | ---------- | | 4th | 1 | 1 | ✓ equal | | 3rd | 4 | 1 | ✓ one is 1 | | 2nd | 2 | 3 | ✘ mismatch | | 1st | 5 | 1 | ✓ one is 1 | ✘ Dimension 2: `2 ≠ 3` and neither is 1 → **broadcasting not allowed**| Index | Name | Description |
|---|---|---|
| 1 | zeros | Allocates initialized memory for a tensor |
| 2 | empty | Allocates uninitialized memory for a tensor |
| 3 | host_to_device | Transfers tensor data from host to device memory |
| 4 | device_to_host | Transfers tensor data from device to host memory |
| 5 | contiguous | Ensures tensor has contiguous memory layout; makes copy if needed. |
Introduced in v0.1.0
Renamed in v1.0.0 (formerly allocate_on_hardware)
Tensor shape - list of dimension sizes. Assumes values are valid and > 0.
This parameter defines the tensor shape, not the data type. The data type is set by the template parameter T
What does host_to_device do? (Click to expand)
The function copies data from a tensor located in host (CPU) memory to a device-specific tensor allocated on accelerator hardware (e.g., GPU, ASIC, FPGA). The actual type of host and device tensors is backend-dependent:
* The **host tensor** may be any type that exposes shape and raw data access (e.g., PyTorch, NumPy).
* The **device tensor** is created as an instance of the backend-defined `DeviceTensorHost-side tensor with shape and data accessible for transfer
The data type must match template type T.
// Create a PyTorch (illustration only) tensor with int32 data on the host (CPU)
torch::Tensor host_tensor = torch::tensor({1, 2, 3}, torch::kInt32);
// Transfer the tensor to device memory using the HEAL interface
auto device_tensor = host_to_device<int32_t>(host_tensor);
What does device_to_hostdo? (Click to expand)
The function enables the copying of data from a tensor residing in device memory (e.g., GPU, FPGA) back to a tensor in host memory (e.g., CPU). This operation is essential for accessing and utilizing the results of computations performed on accelerator hardware within the host application.
The function performs:
1. Allocate a new host tensor of the same shape and type.
2. Copy the device tensor's data to host memory.
3. Return the newly allocated host tensor containing the copied data.
Device-side tensor to be copied to host memory
Data type must match template T.
Why do we need contiguous? (Click to expand)
Some tensor operations, like transposing or slicing, can change the **memory layout** of a tensor without changing its shape.\
This can make the tensor **non-contiguous**, meaning the elements are not laid out sequentially in memory.
A **contiguous tensor** has elements stored in standard row-major order, without gaps or jumps in memory.
Hardware accelerators and many algorithms expect **contiguous tensors** for best performance (and sometimes for correctness).
The `contiguous` function checks whether a tensor is already contiguous:
✔ If it is, nothing changes.
✘ If not, it **creates a contiguous copy** and updates the tensor’s memory and strides.
| Index | Name | Short Description |
|---|---|---|
| 6 | pad_single_axis | Appends zeros at the end of a specific axis, expanding shape. |
| 7 | set_const_val | Sets all elements of a tensor to a constant value; in-place, no allocation. |
What does pad_single_axis do? (Click to expand)
The `pad_single_axis` function takes a tensor and **adds zeros at the end of a chosen axis**.
For example:
* Input: `[1, 2, 3]`, pad = 2, axis = 0 → Result: `[1, 2, 3, 0, 0]`
* Input: `[[1, 2, 3], [4, 5, 6]]`, pad = 2, axis = 1 → Result: `[[1, 2, 3, 0, 0], [4, 5, 6, 0, 0]]`
It doesn’t change the other dimensions — only the one you specify.
Negative axis values count from the end:
* `axis = -1` → last axis,
* `axis = -2` → second-to-last, etc.
**Summary:**\
Copy existing data → Add zeros at end → Output padded tensor.
// This sets every element in the tensor to zero.
auto hw_tensor = host_to_device<int32_t>(torch::rand({10}));
set_const_val<int32_t>(hw_tensor, 0);
| Index | Name | Short Description |
|---|---|---|
| 8 | modmul_ttt | Modular multiplication (tensor-tensor-tensor) |
| 9 | modmul_ttc | Modular multiplication (tensor-tensor-constant) |
| 10 | modmul_tct | Modular multiplication (tensor-constant-tensor) |
| 11 | modmul_tcc | Modular multiplication (tensor-constant-constant) |
| Index | Name | Short Description |
|---|---|---|
| 12 | modsum_ttt | Modular summation (tensor-tensor-tensor) |
| 13 | modsum_ttc | Modular summation (tensor-tensor-constant) |
| 14 | modsum_tct | Modular summation (tensor-constant-tensor) |
| 15 | modsum_tcc | Modular summation (tensor-constant-constant) |
| Index | Name | Short Description |
|---|---|---|
| 16 | mod_tt | tensor % tensor |
| 17 | mod_tc | tensor % scalar |
| 18 | mod_ct | scalar % tensor |
| Index | Name | Short Description |
|---|---|---|
| 19 | modneg_tt | (- tensor) % tensor |
| 20 | modneg_tc | (- tensor) % scalar |
Modulus tensor.
Broadcastable to a
Broadcasts elementwise.
| | `p_scalar` | `T` | Input (`modneg_tc`) | Scalar modulus. | | `result` | `DeviceTensor| Index | Name | Short Description |
|---|---|---|
| 21 | axis_modsum | sums values along a given axis and reduces them modulo p |
| 22 | modmul_axis_sum | Computes a modular sum of products over a specified axis between tensors a and b, with optional permutation |
What does axis_modsum do? (Click to expand)
For each "slice" of the tensor along the selected axis:
* It adds all values in that slice
* Then applies `% p[i]` for each element in the last dimension
* The result is written into the output tensor
#### Example
Assume you have an input tensor `a` with shape `[2, 3, 4]`:
```
a = [
[[ 1, 2, 3, 4],
[ 5, 6, 7, 8],
[ 9, 10, 11, 12]],
[[13, 14, 15, 16],
[17, 18, 19, 20],
[21, 22, 23, 24]]
]
```
And a modulus tensor `p = [11, 13, 17, 19]` (shape `[4]`, matching the last dimension).
Calling:
```cpp
axis_modsum(a, p, result, axis=1);
```
Will reduce across axis 1 (i.e., over the second dimension — the rows). The output will be:
```
[
[(1+5+9)%11, (2+6+10)%13, (3+7+11)%17, (4+8+12)%19],
[(13+17+21)%11, (14+18+22)%13, ...]
]
```
The output shape is `[2, 4]`, same as the input shape with axis 1 removed.
#### Summary
* The reduction is performed across the selected axis
* The modulus is applied element-wise across the last dimension
* The result tensor has one fewer dimension than the input
What does modmul_axis_sum do? (click to expand)
The `modmul_axis_sum` function **multiplies two tensors together along one axis, sums the results, and applies modular reduction**.
In simpler words:
* For each output position, it:
1. Multiplies matching elements from `a` and `b`.
2. Adds up all those multiplied values along a chosen axis.
3. Applies a modulo (`% p`) operation to keep the result within `[0, p)`.
It’s like doing:
```
sum over i (a[...] * b[...]) % p
```
#### Example
```
a = [ [1, 2, 3], [4, 5, 6] ] # shape [2, 3]
b = [ [7, 8, 9], [10, 11, 12] ] # shape [2, 3]
p = [13] # modulus
```
If we sum over axis 1 (columns), we do:
```
For row 0: (1*7 + 2*8 + 3*9) % 13
For row 1: (4*10 + 5*11 + 6*12) % 13
```
The result is a 1D tensor with one value per row, each reduced modulo 13.
***
#### What does `apply_perm` do?
If `apply_perm = true`, you **reorder** the positions along that axis **before multiplying and summing**, using a `perm` vector.
Example:
* If `perm = [2, 0, 1]`, it means:
* “Position 0 becomes 2, 1 becomes 0, 2 becomes 1”
* You apply this shuffle before doing the calculations.
Previously indicated which axis (-1 or -3) represented transform dimension m. |
|
Now the operation always uses axis = -1.
| Index | Name | Short Description |
|---|---|---|
| 23 | ntt | Applies Number-Theoretic Transform - performs a forward transform on batched, multi-channel input tensors, converting them to the NTT domain for efficient polynomial multiplication. |
| 24 | intt | Applies Inverse Number-Theoretic Transform, returning the data to its original (coefficient) domain. |
| Index | Group | Name | Short Description |
|---|---|---|---|
| 25 | Other Compute Operations | apply_g_decomp | Applies gadget decomposition (HE-specific operation) |
| 26 | Other Compute Operations | take_along_axis | Selects values from a tensor along a specified axis using provided indices. |
What does apply_g_decomp_relative_to_full_q do? (Click to expand)
* Input tensor `a` of shape `[reps_l, q_list, reps_r]` contains residues in RNS/CRT form.
* The function reconstructs each element modulo Q with the Q is a product of q\_i (`[0, ..., q_list-1]` )
* Each element is expressed as a sum of digits in base $$2^{g\_base\_bits}$$ : n
$$x = d\_0 \cdot 2^{0 \cdot g\_base\_bits} ;+; d\_1 \cdot 2^{1 \cdot g\_base\_bits} ;+; \dots ;+; d\_{g\_exp-1} \cdot 2^{(g\_exp-1) \cdot g\_base\_bits}$$
* The result is written into a new tensor `out` with shape `[reps_l, g_exp, reps_r]`,\
where the second axis enumerates the digits of the decomposition.
What does take_along_axis do? (Click to expand)
This function allows you to select specific values from a tensor by specifying the exact positions (indices) you want along a chosen axis.
It’s like requesting *“From each row (or column, or depth slice), give me the element at position X.”*
***
### **1D Example**
Suppose you have a 1D tensor:
```
a = [5, 10, 15, 20]
```
And you want to select items in this order: `[2, 0, 3, 1]` → which means:\
pick the 3rd, then 1st, then 4th, then 2nd element.
You call:
```cpp
take_along_axis(a, indices = [2, 0, 3, 1], axis = 0)
```
You get:
```
[15, 5, 20, 10]
```
***
### 2D example:
If you have a 2D tensor:
```
a = [[10, 20, 30],
[40, 50, 60],
[70, 80, 90]]
```
And you want to select:
```
indices = [[2, 1, 0],
[1, 0, 2],
[0, 2, 1]]
```
#### 🔹 Case 1: `axis = 0` → Gather **down rows**
* At each column, we choose from the **vertical direction (axis 0),** so rows are selected.
* Shape of `indices` must match `a`, and each index points to a row.
We apply:
```cpp
take_along_axis(a, indices, axis = 0)
```
Result:
```
[[70, 50, 30], // from rows [2,1,0]
[40, 20, 90], // from rows [1,0,2]
[70, 80, 60]] // from rows [0,2,1]
```
Each column's value is selected from the specified row in that column.
#### 🔹 Case 2: `axis = 1` → Gather **across columns**
* At each row, we choose from the **horizontal direction (axis 1)** — so columns are selected.
We apply:
```cpp
take_along_axis(a, indices, axis = 1)
```
```
[[30, 20, 10], // from columns [2,1,0]
[50, 40, 60], // from columns [1,0,2]
[70, 90, 80]] // from columns [0,2,1]
```
Each value is selected from the **same row**, just pulling the column specified by `indices`.
***
### **Another example**
Input tensor's shape = `[3, 4]` with 3 rows (axis 0) and 4 columns (axis 1):
```
a = [[10, 20, 30, 40],
[50, 60, 70, 80],
[90,100,110,120]]
```
We want to **gather 2 values from each row** →
```
indices = [[2, 0],
[1, 3],
[3, 1]]
```
```cpp
take_along_axis(a, indices, axis=1)
```
Returns:
```
[[ 30, 10],
[ 60, 80],
[120, 100]]
```
The result shape matches `indices.shape`, because the **indices are driving the size along the axis** you’re gathering from.
***
| Index | Name | Short Description |
|---|---|---|
| 27 | flatten | Collapses a range of dimensions into a single axis, reshaping the tensor while preserving element order. |
| 28 | expand | Expands tensor dimensions without copying data |
| 29 | unsqueeze | Adds a dimension of size 1 at a specified position |
| 30 | squeeze | Removes dimensions of size 1 from tensor |
| 31 | reshape | Changes tensor shape while preserving data |
| 32 | moveaxis | Updates dims and strides metadata so that the tensor appears to have the same data but with one axis relocated. |
| 33 | get_slice | Produces a zero-copy sliced view using index or range. |
What does flatten do? (click to expand)
The function **collapses several dimensions of a tensor into a single dimension**, without changing the actual data, just how it’s shaped.
You tell it:
* which dimensions to flatten together (from `start_axis` to `end_axis`),
* and it will **replace that range with one combined dimension**.
#### Simple Analogy
Imagine your tensor is a box of LEGO bricks organized by color, shape, and size:
* `[2, 3, 4]` = 2 colors × 3 shapes × 4 sizes = 24 bricks.
If you flatten from axis `0` to `1`, you mix colors and shapes into one group:
* → `[6, 4]` = 6 combined color-shape combos, still 4 sizes each.
***
### Examples
#### **Flatten middle dimensions**
```cpp
// Tensor shape: [2, 3, 4, 5]
flatten(a, 1, 2)
// Result shape: [2, 12, 5]
```
We flatten `[3, 4]` into `12`.
#### **Flatten all dimensions**
```cpp
// Tensor shape: [2, 3, 4, 5]
flatten(a, 0, -1)
// Result shape: [120]
```
All axes are collapsed into one long row.
#### **Flatten with negative indices**
```cpp
// Tensor shape: [4, 5, 6]
flatten(a, -3, -2)
// Result shape: [20, 6]
```
***
#### ⚠️ Important Notes
* **Does not change the data**, just reshapes the view of it.
* **Input must be contiguous** in memory (no transposes before flatten).
* Axes are **inclusive**, so `flatten(a, 1, 3)` flattens 3 axes, not 2.
tensor = flatten<T>( tensor, start_axis, end_axis);