{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "34d13212",
   "metadata": {},
   "source": [
    "# Tensor Operations"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6999361c",
   "metadata": {},
   "source": [
    "In this tutorial, we'll go over what you can do with tensors. Each supported operation will be written out as a function. Then, all of them will be compiled in a loop and executed with a random input to demonstrate their semantics."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "34fc7213",
   "metadata": {},
   "source": [
    "### Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "a62e11a9",
   "metadata": {},
   "outputs": [],
   "source": [
    "import concrete.numpy as hnp\n",
    "import inspect\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6180966a",
   "metadata": {},
   "source": [
    "### Inputset Definition"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "f8de515c",
   "metadata": {},
   "outputs": [],
   "source": [
    "inputset = [np.random.randint(3, 11, size=(3, 2), dtype=np.uint8) for _ in range(10)]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ae02c598",
   "metadata": {},
   "source": [
    "### Supported Operation Definitions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "d7eeb83c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def reshape(x):\n",
    "    return x.reshape((2, 3))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "68510258",
   "metadata": {},
   "outputs": [],
   "source": [
    "def flatten(x):\n",
    "    return x.flatten()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "db8f502b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def index(x):\n",
    "    return x[2, 0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "5e08a6c4",
   "metadata": {},
   "outputs": [],
   "source": [
    "def slice_(x):\n",
    "    return x.flatten()[1:5]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "b807cc5d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def add_scalar(x):\n",
    "    return x + 10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "59471d3a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def add_tensor(x):\n",
    "    return x + np.array([[1, 2], [3, 3], [2, 1]], dtype=np.uint8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "83bf7d53",
   "metadata": {},
   "outputs": [],
   "source": [
    "def add_tensor_broadcasted(x):\n",
    "    return x + np.array([1, 10], dtype=np.uint8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "ff42df0b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def sub_scalar(x):\n",
    "    return x + (-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "0cc14f94",
   "metadata": {},
   "outputs": [],
   "source": [
    "def sub_tensor(x):\n",
    "    return x + (-np.array([[1, 2], [3, 3], [2, 1]], dtype=np.uint8))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "5e83dd23",
   "metadata": {},
   "outputs": [],
   "source": [
    "def sub_tensor_broadcasted(x):\n",
    "    return x + (-np.array([3, 0], dtype=np.uint8))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "9c68c725",
   "metadata": {},
   "outputs": [],
   "source": [
    "def mul_scalar(x):\n",
    "    return x * 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "66d065e0",
   "metadata": {},
   "outputs": [],
   "source": [
    "def mul_tensor(x):\n",
    "    return x * np.array([[1, 2], [3, 3], [2, 1]], dtype=np.uint8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "a04ae50b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def mul_tensor_broadcasted(x):\n",
    "    return x * np.array([2, 3], dtype=np.uint8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "39fb823b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def power(x):\n",
    "    return x ** 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "4257c1c9",
   "metadata": {},
   "outputs": [],
   "source": [
    "def truediv(x):\n",
    "    return x // 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "712b965a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def dot(x):\n",
    "    return x.flatten() @ np.array([1, 1, 1, 2, 1, 1], dtype=np.uint8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "480b6cc7",
   "metadata": {},
   "outputs": [],
   "source": [
    "def matmul(x):\n",
    "    return x @ np.array([[1, 2, 3], [3, 2, 1]], dtype=np.uint8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "b876272b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def clip(x):\n",
    "    return x.clip(6, 11)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "cec1d224",
   "metadata": {},
   "outputs": [],
   "source": [
    "def comparison(x):\n",
    "    return x > np.array([[10, 5], [8, 11], [3, 7]], dtype=np.uint8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "668ab894",
   "metadata": {},
   "outputs": [],
   "source": [
    "def minimum(x):\n",
    "    return np.minimum(x, np.array([[10, 5], [8, 11], [3, 7]], dtype=np.uint8))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "14031662",
   "metadata": {},
   "outputs": [],
   "source": [
    "def maximum(x):\n",
    "    return np.maximum(x, np.array([[10, 5], [8, 11], [3, 7]], dtype=np.uint8))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12332a5b",
   "metadata": {},
   "source": [
    "Other than these, we support a lot of numpy functions which you can find more about at [Numpy Support](../howto/numpy_support.md)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "09311480",
   "metadata": {},
   "source": [
    "### Compilation and Homomorphic Evaluation of Supported Operations"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cf0152a2",
   "metadata": {},
   "source": [
    "Note that some operations require programmable bootstrapping to work and programmable bootstrapping has a certain probability of failure. Usually, it has more than a 99% probability of success but with big bit-widths, this probability can drop to 95%."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "0cdbc545",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "#######################################################################################\n",
      "\n",
      "def reshape(x):\n",
      "    return x.reshape((2, 3))\n",
      "\n",
      "reshape([[7, 10], [9, 10], [7, 9]]) homomorphically evaluates to [[7, 10, 9], [10, 7, 9]]\n",
      "\n",
      "#######################################################################################\n",
      "\n",
      "def flatten(x):\n",
      "    return x.flatten()\n",
      "\n",
      "flatten([[4, 7], [6, 3], [9, 8]]) homomorphically evaluates to [4, 7, 6, 3, 9, 8]\n",
      "\n",
      "#######################################################################################\n",
      "\n",
      "def index(x):\n",
      "    return x[2, 0]\n",
      "\n",
      "index([[8, 4], [9, 6], [9, 10]]) homomorphically evaluates to 9\n",
      "\n",
      "#######################################################################################\n",
      "\n",
      "def slice_(x):\n",
      "    return x.flatten()[1:5]\n",
      "\n",
      "slice_([[5, 9], [7, 4], [7, 3]]) homomorphically evaluates to [9, 7, 4, 7]\n",
      "\n",
      "#######################################################################################\n",
      "\n",
      "def add_scalar(x):\n",
      "    return x + 10\n",
      "\n",
      "add_scalar([[10, 9], [9, 9], [8, 9]]) homomorphically evaluates to [[20, 19], [19, 19], [18, 19]]\n",
      "\n",
      "#######################################################################################\n",
      "\n",
      "def add_tensor(x):\n",
      "    return x + np.array([[1, 2], [3, 3], [2, 1]], dtype=np.uint8)\n",
      "\n",
      "add_tensor([[6, 9], [9, 7], [7, 5]]) homomorphically evaluates to [[7, 11], [12, 10], [9, 6]]\n",
      "\n",
      "#######################################################################################\n",
      "\n",
      "def add_tensor_broadcasted(x):\n",
      "    return x + np.array([1, 10], dtype=np.uint8)\n",
      "\n",
      "add_tensor_broadcasted([[3, 3], [4, 7], [9, 10]]) homomorphically evaluates to [[4, 13], [5, 17], [10, 20]]\n",
      "\n",
      "#######################################################################################\n",
      "\n",
      "def sub_scalar(x):\n",
      "    return x + (-1)\n",
      "\n",
      "sub_scalar([[3, 8], [10, 4], [5, 8]]) homomorphically evaluates to [[2, 7], [9, 3], [4, 7]]\n",
      "\n",
      "#######################################################################################\n",
      "\n",
      "def sub_tensor(x):\n",
      "    return x + (-np.array([[1, 2], [3, 3], [2, 1]], dtype=np.uint8))\n",
      "\n",
      "sub_tensor([[4, 6], [6, 8], [4, 3]]) homomorphically evaluates to [[3, 4], [3, 5], [2, 2]]\n",
      "\n",
      "#######################################################################################\n",
      "\n",
      "def sub_tensor_broadcasted(x):\n",
      "    return x + (-np.array([3, 0], dtype=np.uint8))\n",
      "\n",
      "sub_tensor_broadcasted([[6, 7], [9, 10], [6, 3]]) homomorphically evaluates to [[3, 7], [6, 10], [3, 3]]\n",
      "\n",
      "#######################################################################################\n",
      "\n",
      "def mul_scalar(x):\n",
      "    return x * 2\n",
      "\n",
      "mul_scalar([[8, 5], [10, 4], [10, 4]]) homomorphically evaluates to [[16, 10], [20, 8], [20, 8]]\n",
      "\n",
      "#######################################################################################\n",
      "\n",
      "def mul_tensor(x):\n",
      "    return x * np.array([[1, 2], [3, 3], [2, 1]], dtype=np.uint8)\n",
      "\n",
      "mul_tensor([[7, 3], [10, 3], [6, 10]]) homomorphically evaluates to [[7, 6], [30, 9], [12, 10]]\n",
      "\n",
      "#######################################################################################\n",
      "\n",
      "def mul_tensor_broadcasted(x):\n",
      "    return x * np.array([2, 3], dtype=np.uint8)\n",
      "\n",
      "mul_tensor_broadcasted([[6, 7], [3, 4], [7, 6]]) homomorphically evaluates to [[12, 21], [6, 12], [14, 18]]\n",
      "\n",
      "#######################################################################################\n",
      "\n",
      "def power(x):\n",
      "    return x ** 2\n",
      "\n",
      "power([[6, 4], [6, 8], [3, 10]]) homomorphically evaluates to [[36, 16], [36, 64], [9, 100]]\n",
      "\n",
      "#######################################################################################\n",
      "\n",
      "def truediv(x):\n",
      "    return x // 2\n",
      "\n",
      "truediv([[9, 5], [4, 7], [3, 7]]) homomorphically evaluates to [[4, 2], [2, 3], [1, 3]]\n",
      "\n",
      "#######################################################################################\n",
      "\n",
      "def dot(x):\n",
      "    return x.flatten() @ np.array([1, 1, 1, 2, 1, 1], dtype=np.uint8)\n",
      "\n",
      "dot([[4, 10], [4, 3], [5, 8]]) homomorphically evaluates to 37\n",
      "\n",
      "#######################################################################################\n",
      "\n",
      "def matmul(x):\n",
      "    return x @ np.array([[1, 2, 3], [3, 2, 1]], dtype=np.uint8)\n",
      "\n",
      "matmul([[7, 10], [3, 8], [3, 3]]) homomorphically evaluates to [[37, 34, 31], [27, 22, 17], [12, 12, 12]]\n",
      "\n",
      "#######################################################################################\n",
      "\n",
      "def clip(x):\n",
      "    return x.clip(6, 11)\n",
      "\n",
      "clip([[9, 5], [10, 10], [4, 6]]) homomorphically evaluates to [[9, 6], [10, 10], [6, 6]]\n",
      "\n",
      "#######################################################################################\n",
      "\n",
      "def comparison(x):\n",
      "    return x > np.array([[10, 5], [8, 11], [3, 7]], dtype=np.uint8)\n",
      "\n",
      "comparison([[4, 4], [5, 10], [7, 6]]) homomorphically evaluates to [[0, 0], [0, 0], [1, 0]]\n",
      "\n",
      "#######################################################################################\n",
      "\n",
      "def maximum(x):\n",
      "    return np.maximum(x, np.array([[10, 5], [8, 11], [3, 7]], dtype=np.uint8))\n",
      "\n",
      "maximum([[9, 7], [4, 10], [6, 8]]) homomorphically evaluates to [[10, 7], [8, 11], [6, 8]]\n",
      "\n",
      "#######################################################################################\n",
      "\n",
      "def minimum(x):\n",
      "    return np.minimum(x, np.array([[10, 5], [8, 11], [3, 7]], dtype=np.uint8))\n",
      "\n",
      "minimum([[4, 9], [10, 8], [10, 7]]) homomorphically evaluates to [[4, 5], [8, 8], [3, 7]]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "functions = [\n",
    "    reshape,\n",
    "    flatten,\n",
    "    index,\n",
    "    slice_,\n",
    "    add_scalar,\n",
    "    add_tensor,\n",
    "    add_tensor_broadcasted,\n",
    "    sub_scalar,\n",
    "    sub_tensor,\n",
    "    sub_tensor_broadcasted,\n",
    "    mul_scalar,\n",
    "    mul_tensor,\n",
    "    mul_tensor_broadcasted,\n",
    "    power,\n",
    "    truediv,\n",
    "    dot,\n",
    "    matmul,\n",
    "    clip,\n",
    "    comparison,\n",
    "    maximum,\n",
    "    minimum,\n",
    "]\n",
    "\n",
    "for function in functions:\n",
    "    compiler = hnp.NPFHECompiler(function, {\"x\": \"encrypted\"})\n",
    "    circuit = compiler.compile_on_inputset(inputset)\n",
    "    \n",
    "    sample = np.random.randint(3, 11, size=(3, 2), dtype=np.uint8)\n",
    "    result = circuit.run(sample)\n",
    "    \n",
    "    print(\"#######################################################################################\")\n",
    "    \n",
    "    print()\n",
    "    print(f\"{inspect.getsource(function)}\")\n",
    "    print(f\"{function.__name__}({sample.tolist()}) homomorphically evaluates to {result if isinstance(result, int) else result.tolist()}\")\n",
    "    print()\n",
    "\n",
    "    expected = function(sample)\n",
    "    if not np.array_equal(result, expected):\n",
    "        print(f\"(It should have been evaluated to {expected if isinstance(expected, int) else expected.tolist()} but it didn't due to an error during PBS)\")\n",
    "        print()"
   ]
  }
 ],
 "metadata": {
  "execution": {
   "timeout": 10800
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}