Tensor operations

The module std::tensor gives a tensor the operations of numerical code. It has sums, statistics, element-wise changes, reductions along an axis and a matrix product. A tour of Vx shows how to make a tensor, read its shape, and reshape or pad it. This chapter starts from there.

Each program in this chapter is also a test of the compiler. You can find it under tests/backend/pass/ with the name that its section gives.

Sums of a 1-D tensor

A reduction turns many elements into one value. A sum is the most common one. Add import std::tensor; to get these methods on a 1-D tensor of f16, bf16, f32 or f64:

• a.sum() adds the elements, and a.mean() divides that sum by the length.
• a.dot(&b) multiplies the elements of a and b in pairs and adds the products.
• a.norm() is the length of a as a vector: the square root of a.dot(&a).
• a.max() and a.min() give the largest and the smallest element. They skip a NaN.

This program calls each one (book_tensor_ops_sums.vx):

import std::tensor;

fn main() -> i32 {
  let a : Tensor<f32, [4]> = [ 1.0, -2.0, 3.0, 2.0 ];
  let b : Tensor<f32, [4]> = [ 2.0, 1.0, 0.5, -1.0 ];
  println!("sum {} dot {}", a.sum(), a.dot(&b));
  println!("max {} min {}", a.max(), a.min());
  println!("mean {} norm {}", a.mean(), a.norm());
  return 0;
}

It prints sum 4 dot -0.5, then max 3 min -2, then mean 1 norm 4.2426405.

The length can be fixed in the type, as in Tensor<f32, [4]>, or known only at run time, as in Tensor<f32, [?]>. A row of a 2-D tensor is a 1-D tensor too, so m[1].sum() adds row 1.

These methods add the elements in an order that they choose, so that they can use vector instructions. The answer can differ in the last bits from a plain loop that adds one element at a time. The methods add f16 or bf16 elements in f32, and they give an f32. An f16 total above 65504 does not fit in an f16.

In earlier versions of Vx, the compiler had dot, sum, max and min as built-in functions. Now they are methods. So dot(q[i], k[j]) is now q[i].dot(&k[j]), and the program needs import std::tensor;.

Statistics and softmax

std::tensor has these methods on a 1-D tensor too:

• a.argmax() and a.argmin() give the index of the largest or smallest element. They give an Option<i32>: None for an empty tensor. When two elements are equal, they give the first one.
• a.variance() is the mean of the squared distances from the mean. It divides by the length n, not n - 1. a.std_dev() is its square root.
• a.softmax_inplace() replaces each element x with exp(x) divided by the sum of exp over all the elements. After it, the elements are positive and their sum is 1.

softmax_inplace subtracts the largest element before it calls exp. The answer is the same, but exp(1000.0) does not overflow. This program uses each method (book_tensor_ops_statistics.vx):

import std::tensor;

fn main() -> i32 {
  let scores : Tensor<f64, [5]> = [ 2.0, 4.0, 4.0, 1.0, 9.0 ];
  match scores.argmax() {
    Option<i32>::Some(i) => println!("argmax {}", i),
    Option<i32>::None => println!("argmax none"),
  }
  println!("argmin {}", scores.argmin().unwrap());
  println!("variance {} std_dev {}", scores.variance(), scores.std_dev());

  let mut logits : Tensor<f64, [3]> = [ 1000.0, 1001.0, 1002.0 ];
  logits.softmax_inplace();
  println!("softmax {} {} {}", logits[0], logits[1], logits[2]);

  let empty = Tensor<f64>([0]);
  println!("empty argmax is none: {}", empty.argmax().is_none());
  return 0;
}

It prints argmax 4, argmin 3, variance 7.6 std_dev 2.756809750418044, then softmax 0.09003057317038046 0.24472847105479764 0.6652409557748218 and empty argmax is none: true. The last digits of the softmax line can differ on another machine, because exp comes from the system’s math library.

softmax_inplace changes the tensor, so the tensor must be mut. The other methods of this section do not change the tensor.

Element-wise operations in place

An element-wise operation does the same thing to each element. The methods of std::tensor change the tensor in place, so they make no new tensor. A chain of them allocates no memory.

The _assign methods take the place of += and the other compound operators, because Vx cannot overload them on a tensor yet. Both tensors must have the same length. A different length stops the program with an error message. This program uses them one after the other (book_tensor_ops_in_place.vx):

import std::tensor;

fn show(t : &Tensor<f32, [?]>) -> i32 {
  for i in 0..t.len() {
    print!(" {}", t[i]);
  }
  println!("");
  return 0;
}

fn main() -> i32 {
  let mut y : Tensor<f32, [4]> = [ 1.0, 2.0, 3.0, 4.0 ];
  let x : Tensor<f32, [4]> = [ 10.0, 20.0, 30.0, 40.0 ];
  y.add_assign(&x);
  print!("add_assign");
  show(&y);
  y.scale(0.5);
  print!("scale");
  show(&y);
  y.scaled_add(-1.0, &x);
  print!("scaled_add");
  show(&y);
  y.clamp_inplace(-10.0, 0.0);
  print!("clamp_inplace");
  show(&y);
  y.map_inplace(| v : f32 | v * v);
  print!("map_inplace");
  show(&y);
  y.sqrt_inplace();
  print!("sqrt_inplace");
  show(&y);
  return 0;
}

It prints one line for each step, from add_assign 11 22 33 44 to sqrt_inplace 4.5 9 10 10.

map_inplace calls the closure once for each element. When one of the named methods does the same work, use it. It is faster.

Your own reduction with reduce

t.reduce(init, |acc, x| ..) is a reduction that you write. It starts with acc equal to init. For each element x, it calls the closure and makes the result the new acc. At the end, it gives the last acc. reduce is part of the language, so it needs no import. It works on a tensor of any rank and of any element type, integers too. The closure can read a variable from outside it.

This program finds a total, a largest value, a total of reduced prices, and a product (book_tensor_ops_reduce.vx):

fn main() -> i32 {
  let prices : Tensor<f64, [4]> = [ 3.0, 7.5, 2.0, 5.5 ];
  let total = prices.reduce(0.0, | acc, x | acc + x);
  let most = prices.reduce(0.0, | acc, x | if x > acc {
    x
  } else {
    acc
  });
  let discount : f64 = 0.5;
  let sale = prices.reduce(0.0, | acc, x | acc + x * discount);
  println!("total {} most {} sale {}", total, most, sale);

  let counts : Tensor<i32, [2, 3]> = [ [ 1, 2, 3 ], [ 4, 5, 6 ] ];
  println!("product {}", counts.reduce(1, | acc, x | acc * x));
  return 0;
}

It prints total 18 most 7.5 sale 9, then product 720.

init must have the type of the elements, and the closure must give that type too. The compiler refuses a program that breaks either rule. reduce can combine the elements in any order, as sum does. So use a closure for which the order does not change the answer, such as +, * or a maximum.

Matrices: reductions along an axis and matmul

A 2-D tensor of a float type has sum, mean, max and min over all its elements. It has the element-wise methods of the table above, except scaled_add and map_inplace. It also has reductions along an axis, as numpy’s axis= argument does:

• m.sum_axis(0) reduces each column, so it gives one value for each column.
• m.sum_axis(1) reduces each row, so it gives one value for each row.
• mean_axis, max_axis and min_axis work in the same way.

Each one gives a new 1-D tensor. An axis other than 0 or 1 stops the program.

a.matmul(&b) is the matrix product for f32 and f64. It makes a new tensor with the rows of a and the columns of b. Element [i][j] is the sum over k of a[i][k] * b[k][j]. This program uses both (book_tensor_ops_matrices.vx):

import std::tensor;

fn show(t : &Tensor<f64, [?]>) -> i32 {
  for i in 0..t.len() {
    print!(" {}", t[i]);
  }
  println!("");
  return 0;
}

fn main() -> i32 {
  let m : Tensor<f64, [2, 3]> = [ [ 1.0, 2.0, 3.0 ], [ 4.0, 5.0, 6.0 ] ];
  println!("sum {} max {}", m.sum(), m.max());
  print!("sum_axis(0)");
  show(&m.sum_axis(0));
  print!("sum_axis(1)");
  show(&m.sum_axis(1));
  print!("mean_axis(1)");
  show(&m.mean_axis(1));

  let k : Tensor<f64, [3, 2]> = [ [ 1.0, 0.0 ], [ 0.0, 1.0 ], [ 1.0, 1.0 ] ];
  let p = m.matmul(&k);
  println!("matmul is {} by {}", p.extent(0), p.extent(1));
  for i in 0..2 {
    println!("{} {}", p[i][0], p[i][1]);
  }
  return 0;
}

It prints sum 21 max 6, sum_axis(0) 5 7 9, sum_axis(1) 6 15 and mean_axis(1) 2 5. Then it prints matmul is 2 by 2 and the two rows 4 5 and 10 11.

The number of columns of a must be equal to the number of rows of b. matmul checks this when the program runs, not when it compiles, even when the two shapes are in the types. A wrong shape stops the program with the message matmul needs as many columns on the left as rows on the right.

For code on a GPU, see Parallel code on a GPU.