Data Types
==========

Every tensor has a data type, which is typically ``float32`` in deep
learning, but also could be ``int8`` (e.g. for model quantization) and
others. The ``tvm_vector_add`` module we created in
:numref:`ch_vector_add` only accepts ``float32`` tensors. Let's extend
it to other data types in this section.

Specifying a Data Type
----------------------

To use a data type different to the default ``float32``, we can specify
it explicitly when creating the placeholders. In the following code
block, we generalize the vector addition expression defined in
:numref:`ch_vector_add` to accept an argument ``dtype`` to specify the
data type. In particular, we pass ``dtype`` to ``te.placeholder`` when
creating ``A`` and ``B``. The result ``C`` then obtains the same data
type as ``A`` and ``B``.

.. code:: python

    import tvm
    from tvm import te
    import d2ltvm
    import numpy as np
    
    n = 100
    
    def tvm_vector_add(dtype):
        A = te.placeholder((n,), dtype=dtype)
        B = te.placeholder((n,), dtype=dtype)
        C = te.compute(A.shape, lambda i: A[i] + B[i])
        print('expression dtype:', A.dtype, B.dtype, C.dtype)
        s = te.create_schedule(C.op)
        return tvm.build(s, [A, B, C])

Let's compile a module that accepts ``int32`` tensors.

.. code:: python

    mod = tvm_vector_add('int32')


.. parsed-literal::
    :class: output

    expression dtype: int32 int32 int32


Then we define a method to verify the results with a particular data
type. Note that we pass a constructor that modifies the tensor data type
by ``astype``.

.. code:: python

    def test_mod(mod, dtype):
        a, b, c = d2ltvm.get_abc(n, lambda x: tvm.nd.array(x.astype(dtype)))
        print('tensor dtype:', a.dtype, b.dtype, c.dtype)
        mod(a, b, c)
        np.testing.assert_equal(c.asnumpy(), a.asnumpy() + b.asnumpy())
    
    test_mod(mod, 'int32')


.. parsed-literal::
    :class: output

    tensor dtype: int32 int32 int32


Let's try other data types as well

.. code:: python

    for dtype in ['float16', 'float64', 'int8','int16', 'int64']:
        mod = tvm_vector_add(dtype)
        test_mod(mod, dtype)


.. parsed-literal::
    :class: output

    expression dtype: float16 float16 float16
    tensor dtype: float16 float16 float16
    expression dtype: float64 float64 float64
    tensor dtype: float64 float64 float64
    expression dtype: int8 int8 int8
    tensor dtype: int8 int8 int8
    expression dtype: int16 int16 int16
    tensor dtype: int16 int16 int16
    expression dtype: int64 int64 int64
    tensor dtype: int64 int64 int64


Converting Elements Data Types
------------------------------

Besides constructing a tensor with a particular data type, we can also
cast the data type of a tensor element during the computation. The
following method is the same as ``tvm_vector_add`` except that it casts
the data type of A and B in ``te.compute``, leaving the data type
defined in ``te.placeholder`` as default (``float32``). Because of the
casting done by ``astype``, the result ``C`` will have the data type
specified by ``dtype``.

.. code:: python

    def tvm_vector_add_2(dtype):
        A = te.placeholder((n,))
        B = te.placeholder((n,))
        C = te.compute(A.shape, 
                        lambda i: A[i].astype(dtype) + B[i].astype(dtype))
        print('expression dtype:', A.dtype, B.dtype, C.dtype)
        s = te.create_schedule(C.op)
        return tvm.build(s, [A, B, C])

Then we define a similar test function to verify the results.

.. code:: python

    def test_mod_2(mod, dtype):
        a, b, c = d2ltvm.get_abc(n)
        # by default `get_abc` returns NumPy ndarray in float32
        a_tvm, b_tvm = tvm.nd.array(a), tvm.nd.array(b)
        c_tvm = tvm.nd.array(c.astype(dtype))
        print('tensor dtype:', a_tvm.dtype, b_tvm.dtype, c_tvm.dtype)
        mod(a_tvm, b_tvm, c_tvm)
        np.testing.assert_equal(c_tvm.asnumpy(), a.astype(dtype) + b.astype(dtype))
    
    mod = tvm_vector_add_2('int32')
    test_mod_2(mod, 'int32')


.. parsed-literal::
    :class: output

    expression dtype: float32 float32 int32
    tensor dtype: float32 float32 int32


Summary
-------

-  We can specify the data type by ``dtype`` when creating TVM
   placeholders.
-  The data type of a tensor element can be cast by ``astype`` in TVM
   compute.