{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "34d13212",
   "metadata": {},
   "source": [
    "# Working With Tensors"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6999361c",
   "metadata": {},
   "source": [
    "In this tutorial, we'll go over what you can do with encrypted 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": "markdown",
   "id": "ab71e23f",
   "metadata": {},
   "source": [
    "We will generate some random input tensors as calibration data for our encrypted tensor functions."
   ]
  },
  {
   "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": "e917b82a",
   "metadata": {},
   "source": [
    "### Prepare Supported Operations List "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9495a29d",
   "metadata": {},
   "source": [
    "We will create a list of supported operations to showcase them in a loop."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "0cb14b31",
   "metadata": {},
   "outputs": [],
   "source": [
    "supported_operations = [\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",
    "]"
   ]
  },
  {
   "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": 25,
   "id": "0cdbc545",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "#######################################################################################\n",
      "\n",
      "def reshape(x):\n",
      "    return x.reshape((2, 3))\n",
      "\n",
      "reshape([[3, 6], [5, 6], [9, 10]]) homomorphically evaluates to [[3, 6, 5], [6, 9, 10]]\n",
      "\n",
      "#######################################################################################\n",
      "\n",
      "def flatten(x):\n",
      "    return x.flatten()\n",
      "\n",
      "flatten([[7, 8], [10, 9], [8, 9]]) homomorphically evaluates to [7, 8, 10, 9, 8, 9]\n",
      "\n",
      "#######################################################################################\n",
      "\n",
      "def index(x):\n",
      "    return x[2, 0]\n",
      "\n",
      "index([[3, 10], [5, 4], [6, 4]]) homomorphically evaluates to 6\n",
      "\n",
      "#######################################################################################\n",
      "\n",
      "def slice_(x):\n",
      "    return x.flatten()[1:5]\n",
      "\n",
      "slice_([[5, 7], [5, 6], [9, 5]]) homomorphically evaluates to [7, 5, 6, 9]\n",
      "\n",
      "#######################################################################################\n",
      "\n",
      "def add_scalar(x):\n",
      "    return x + 10\n",
      "\n",
      "add_scalar([[3, 5], [4, 8], [9, 5]]) homomorphically evaluates to [[13, 15], [14, 18], [19, 15]]\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([[4, 3], [4, 9], [8, 3]]) homomorphically evaluates to [[5, 5], [7, 12], [10, 4]]\n",
      "\n",
      "#######################################################################################\n",
      "\n",
      "def add_tensor_broadcasted(x):\n",
      "    return x + np.array([1, 10], dtype=np.uint8)\n",
      "\n",
      "add_tensor_broadcasted([[9, 3], [4, 4], [8, 6]]) homomorphically evaluates to [[10, 13], [5, 14], [9, 16]]\n",
      "\n",
      "#######################################################################################\n",
      "\n",
      "def sub_scalar(x):\n",
      "    return x + (-1)\n",
      "\n",
      "sub_scalar([[6, 6], [5, 10], [4, 9]]) homomorphically evaluates to [[5, 5], [4, 9], [3, 8]]\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([[7, 3], [6, 3], [9, 5]]) homomorphically evaluates to [[6, 1], [3, 0], [7, 4]]\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], [10, 6], [3, 10]]) homomorphically evaluates to [[3, 7], [7, 6], [0, 10]]\n",
      "\n",
      "#######################################################################################\n",
      "\n",
      "def mul_scalar(x):\n",
      "    return x * 2\n",
      "\n",
      "mul_scalar([[10, 4], [8, 6], [7, 7]]) homomorphically evaluates to [[20, 8], [16, 12], [14, 14]]\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([[10, 8], [3, 6], [8, 4]]) homomorphically evaluates to [[10, 16], [9, 18], [16, 4]]\n",
      "\n",
      "#######################################################################################\n",
      "\n",
      "def mul_tensor_broadcasted(x):\n",
      "    return x * np.array([2, 3], dtype=np.uint8)\n",
      "\n",
      "mul_tensor_broadcasted([[4, 5], [9, 7], [9, 5]]) homomorphically evaluates to [[8, 15], [18, 21], [18, 15]]\n",
      "\n",
      "#######################################################################################\n",
      "\n",
      "def power(x):\n",
      "    return x ** 2\n",
      "\n",
      "power([[10, 9], [9, 10], [8, 7]]) homomorphically evaluates to [[100, 81], [81, 100], [64, 49]]\n",
      "\n",
      "#######################################################################################\n",
      "\n",
      "def truediv(x):\n",
      "    return x // 2\n",
      "\n",
      "truediv([[10, 7], [7, 7], [4, 8]]) homomorphically evaluates to [[5, 3], [3, 3], [2, 4]]\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([[3, 10], [4, 7], [7, 6]]) homomorphically evaluates to 44\n",
      "\n",
      "#######################################################################################\n",
      "\n",
      "def matmul(x):\n",
      "    return x @ np.array([[1, 2, 3], [3, 2, 1]], dtype=np.uint8)\n",
      "\n",
      "matmul([[8, 9], [5, 5], [8, 9]]) homomorphically evaluates to [[35, 34, 33], [20, 20, 20], [35, 34, 33]]\n",
      "\n",
      "#######################################################################################\n",
      "\n",
      "def clip(x):\n",
      "    return x.clip(6, 11)\n",
      "\n",
      "clip([[3, 4], [4, 4], [8, 7]]) homomorphically evaluates to [[6, 6], [6, 6], [8, 7]]\n",
      "\n",
      "#######################################################################################\n",
      "\n",
      "def comparison(x):\n",
      "    return x > np.array([[10, 5], [8, 11], [3, 7]], dtype=np.uint8)\n",
      "\n",
      "comparison([[3, 5], [8, 8], [3, 7]]) homomorphically evaluates to [[0, 0], [0, 0], [0, 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([[5, 10], [4, 9], [9, 6]]) homomorphically evaluates to [[10, 10], [8, 11], [9, 7]]\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([[9, 8], [4, 3], [5, 9]]) homomorphically evaluates to [[9, 5], [4, 3], [3, 7]]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "for operation in supported_operations:\n",
    "    compiler = hnp.NPFHECompiler(operation, {\"x\": \"encrypted\"})\n",
    "    circuit = compiler.compile_on_inputset(inputset)\n",
    "    \n",
    "    # We setup an example tensor that will be encrypted and passed on to the current operation\n",
    "    sample = np.random.randint(3, 11, size=(3, 2), dtype=np.uint8)\n",
    "    result = circuit.run(sample)\n",
    "    \n",
    "    print(\"#######################################################################################\")\n",
    "    print()\n",
    "    print(f\"{inspect.getsource(operation)}\")\n",
    "    print(f\"{operation.__name__}({sample.tolist()}) homomorphically evaluates to {result if isinstance(result, int) else result.tolist()}\")\n",
    "    print()\n",
    "\n",
    "    expected = operation(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
}