# 2.2. Shapes¶ Open the notebook in Colab

The vector addition module defined in Section 1.2 only accepts vectors with 100-length. It’s too restrictive for real scenarios where inputs can have arbitrary shapes. In this section, we will show how to relax this constraint to deal with general cases.

## 2.2.1. Variable Shapes¶

Remember that we create symbolic placeholders for tensors A and B so we can feed with data later. We can do the same thing for the shape as well. In particular, the following code block uses te.var to create a symbolic variable for an int32 scalar, whose value can be specified later.

import d2ltvm
import numpy as np
import tvm
from tvm import te

n = te.var(name='n')
type(n), n.dtype

(tvm.tir.expr.Var, 'int32')


Now we can use (n,) to create a placeholder for an arbitrary length vector.

A = te.placeholder((n,), name='a')
B = te.placeholder((n,), name='b')
C = te.compute(A.shape, lambda i: A[i] + B[i], name='c')
s = te.create_schedule(C.op)
tvm.lower(s, [A, B, C], simple_mode=True)

produce c {
for (i, 0, n) {
c[(i*stride)] = (a[(i*stride)] + b[(i*stride)])
}
}


Compared to the generated pseudo codes in Section 1.2, we can see the upper bound value of the for loop is changed from 100 to n.

Now we define a similar test function as before to verify that the compiled module is able to correctly execute on input vectors with different lengths.

def test_mod(mod, n):
a, b, c = d2ltvm.get_abc(n, tvm.nd.array)
mod(a, b, c)
print('c.shape:', c.shape)
np.testing.assert_equal(c.asnumpy(), a.asnumpy() + b.asnumpy())

mod = tvm.build(s, [A, B, C])
test_mod(mod, 5)
test_mod(mod, 1000)

c.shape: (5,)
c.shape: (1000,)


But note that we still place the constraint that A, B, and C must be in the same shape. So an error will occur if it is not satisfied.

## 2.2.2. Multi-dimensional Shapes¶

You may already notice that a shape is presented as a tuple. A single element tuple means a 1-D tensor, or a vector. We can extend it to multi-dimensional tensors by adding variables to the shape tuple.

The following method builds a module for multi-dimensional tensor addition, the number of dimensions is specified by ndim. For a 2-D tensor, we can access its element by A[i,j], similarly A[i,j,k] for 3-D tensors. Note that we use *i to handle the general multi-dimensional case in the following code.

def tvm_vector_add(ndim):
A = te.placeholder([te.var() for _ in range(ndim)])
B = te.placeholder(A.shape)
C = te.compute(A.shape, lambda *i: A[i] + B[i])
s = te.create_schedule(C.op)
return tvm.build(s, [A, B, C])


Verify that it works beyond vectors.

mod = tvm_vector_add(2)
test_mod(mod, (2, 2))


c.shape: (2, 2)

• We can use te.var() to specify the dimension(s) of a shape when we don’t know the concrete data shape before execution.
• The shape of an $$n$$-dimensional tensor is presented as an $$n$$-length tuple.