notebooks/differentiable-indexed-arrays.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Differentiable indexing of arrays\n",
    "Since, table or array lookup is an inherently non-differentiable process, the autograd is unable to resolve gradients of the result with respect to the output.\n",
    "\n",
    "Here we have some strategies which gives a valid gradient with respect to both the index and input array. The choice of the strategy is problem specific. Also, this is by no means an exhaustive list. You can help by expanding it.\n",
    "\n",
    "Abbreviation: WRT = With respect to\n",
    "\n",
    "Further Reading: [Neural Turing Machines](https://arxiv.org/abs/1410.5401)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "from library.statistical_math import to_prob_dist_all, entropy"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Naive lookup\n",
    "Naive lookup does produce a gradient wrt its input array but not wrt the index."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": "tf.Tensor(3.0, shape=(), dtype=float32)\ntf.Tensor([0. 0. 1. 0. 0.], shape=(5,), dtype=float32)\nNone\n"
    }
   ],
   "source": [
    "@tf.function\n",
    "def naive_lookup(arr, index):\n",
    "    index = tf.round(index)\n",
    "    index = tf.cast(index, tf.int32)\n",
    "    result = arr[index]\n",
    "    return result\n",
    "\n",
    "arr = tf.Variable([1,2,3,4,5],dtype=tf.float32)\n",
    "index = tf.Variable(1.5, dtype=tf.float32)\n",
    "\n",
    "with tf.GradientTape(persistent=True) as tape:\n",
    "    z = naive_lookup(arr, index)\n",
    "\n",
    "print(z)\n",
    "print(tape.gradient(z, arr))\n",
    "print(tape.gradient(z, index))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Linear lookup\n",
    "In this method we use [linear interpolation](https://en.wikipedia.org/wiki/Linear_interpolation) to interpolate between the two nearest candidates. For 2D arrays, [Bilinear interpolation](https://en.wikipedia.org/wiki/Bilinear_interpolation) can be used.\n",
    "\n",
    "This gives a well defined gradient wrt to both the input and index. However, it is a soft lookup and can return values not present in the array itself.\n",
    "\n",
    "One of the downsides of using this method is that it can only lookup adjacent cells in the [number line](https://en.wikipedia.org/wiki/Number_line).\n",
    "\n",
    "![number line](images/1125px-Number-line.svg.png)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "@tf.function\n",
    "def interp_factor(index):\n",
    "    t1 = tf.math.floor(index)\n",
    "    t2 = tf.math.ceil(index)\n",
    "    \n",
    "    t = tf.math.divide_no_nan((index - t1), (t2 - t1))\n",
    "    \n",
    "    i1 = tf.cast(t1, tf.int32)\n",
    "    i2 = tf.cast(t2, tf.int32)\n",
    "    \n",
    "    return t, i1, i2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": "tf.Tensor(2.5, shape=(), dtype=float32)\ntf.Tensor([0.  0.5 0.5 0.  0. ], shape=(5,), dtype=float32)\ntf.Tensor(-1.0, shape=(), dtype=float32)\n"
    }
   ],
   "source": [
    "@tf.function\n",
    "def linear_lookup(arr, index):\n",
    "    t, i1, i2 = interp_factor(index)\n",
    "    \n",
    "    # Linear interpolation\n",
    "    result = t * arr[i1] + (1 - t) * arr[i2]\n",
    "    \n",
    "    return result\n",
    "\n",
    "arr = tf.Variable([1,2,3,4,5],dtype=tf.float32)\n",
    "index = tf.Variable(1.5, dtype=tf.float32)\n",
    "\n",
    "with tf.GradientTape(persistent=True) as tape:\n",
    "    z = linear_lookup(arr, index)\n",
    "\n",
    "print(z)\n",
    "print(tape.gradient(z, arr))\n",
    "print(tape.gradient(z, index))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Superposition lookup\n",
    "In this method, we have a distribution instead of an integer index. This distribution usually comes after a softmax operation. The result is the dot product of the index and the input array. This is a very popular method in DNN literature."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": "tf.Tensor([0.  0.  0.5 0.  0.5], shape=(5,), dtype=float32)\ntf.Tensor(\n[[0.5 0.5 0.5 0.5 0.5]\n [0.5 0.5 0.5 0.5 0.5]\n [0.  0.  0.  0.  0. ]\n [0.  0.  0.  0.  0. ]\n [0.  0.  0.  0.  0. ]], shape=(5, 5), dtype=float32)\ntf.Tensor([1. 1. 1. 1. 1.], shape=(5,), dtype=float32)\n"
    }
   ],
   "source": [
    "@tf.function\n",
    "def superposition_lookup_vectored(arr, indices):\n",
    "    if tf.rank(arr) == 1:\n",
    "        arr = tf.expand_dims(arr, -1)\n",
    "    indices = tf.expand_dims(indices, -1)\n",
    "    result = arr * indices\n",
    "    return tf.reduce_sum(result, axis=0)\n",
    "\n",
    "arr = tf.Variable([\n",
    "    [0, 0, 1, 0, 0],\n",
    "    [0, 0, 0, 0, 1],\n",
    "    [0, 1, 0, 0, 0],\n",
    "    [1, 0, 0, 0, 0],\n",
    "    [0, 0, 0, 1, 0],\n",
    "],dtype=tf.float32)\n",
    "indices = tf.Variable([0.5, 0.5, 0, 0, 0], dtype=tf.float32)\n",
    "\n",
    "with tf.GradientTape(persistent=True) as tape:\n",
    "#     indices = tf.nn.softmax(indices)\n",
    "    z = superposition_lookup_vectored(arr, indices)\n",
    "\n",
    "print(z)\n",
    "print(tape.gradient(z, arr))\n",
    "print(tape.gradient(z, indices))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": "tf.Tensor([3.0000002], shape=(1,), dtype=float32)\ntf.Tensor([0.  0.1 0.8 0.  0.1], shape=(5,), dtype=float32)\ntf.Tensor([1. 2. 3. 5. 4.], shape=(5,), dtype=float32)\n"
    }
   ],
   "source": [
    "arr = tf.Variable([1,2,3,5,4],dtype=tf.float32)\n",
    "indices = tf.Variable([0.0, 0.1, 0.8, 0.0, 0.1], dtype=tf.float32)\n",
    "\n",
    "with tf.GradientTape(persistent=True) as tape:\n",
    "#     indices = tf.nn.softmax(indices)\n",
    "    z = superposition_lookup_vectored(arr, indices)\n",
    "\n",
    "print(z)\n",
    "print(tape.gradient(z, arr))\n",
    "print(tape.gradient(z, indices))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": "tf.Tensor([0.  0.  0.5 0.5 0. ], shape=(5,), dtype=float32)\ntf.Tensor([2.5], shape=(1,), dtype=float32)\ntf.Tensor(-1.0, shape=(), dtype=float32)\n"
    }
   ],
   "source": [
    "@tf.function\n",
    "def bandwidthify(index, bandwidth):\n",
    "    t, i1, i2 = interp_factor(index)\n",
    "    \n",
    "    # Prevent array out of bounds\n",
    "    i1 = tf.clip_by_value(i1, 0, bandwidth - 1)\n",
    "    i2 = tf.clip_by_value(i2, 0, bandwidth - 1)\n",
    "    t = tf.clip_by_value(t, 0, 1)\n",
    "    \n",
    "    # Linear interpolation\n",
    "    eye = tf.eye(bandwidth)\n",
    "    result = t * eye[i1] + (1 - t) * eye[i2]\n",
    "    \n",
    "    return result\n",
    "\n",
    "index = tf.Variable(2.5, dtype=tf.float32)\n",
    "bandwidth = tf.constant(5, dtype=tf.int32)\n",
    "dummy_array = tf.cast(tf.range(bandwidth), tf.float32)\n",
    "with tf.GradientTape() as tape:\n",
    "    z = bandwidthify(index, bandwidth)\n",
    "    nz = superposition_lookup_vectored(dummy_array, z) # Lookup operation\n",
    "\n",
    "print(z)\n",
    "print(nz)\n",
    "print(tape.gradient(nz, index))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": "tf.Tensor([2.5], shape=(1,), dtype=float32)\ntf.Tensor([0.  0.5 0.5 0.  0. ], shape=(5,), dtype=float32)\ntf.Tensor(-1.0, shape=(), dtype=float32)\n"
    }
   ],
   "source": [
    "@tf.function\n",
    "def superposition_lookup(arr, index):\n",
    "    bandwidth = tf.shape(arr)[0]\n",
    "    vectored_index = bandwidthify(index, bandwidth)\n",
    "    result = superposition_lookup_vectored(arr, vectored_index)\n",
    "    \n",
    "    return result\n",
    "\n",
    "arr = tf.Variable([1,2,3,4,5],dtype=tf.float32)\n",
    "index = tf.Variable(1.5, dtype=tf.float32)\n",
    "\n",
    "with tf.GradientTape(persistent=True) as tape:\n",
    "    z = superposition_lookup(arr, index)\n",
    "\n",
    "print(z)\n",
    "print(tape.gradient(z, arr))\n",
    "print(tape.gradient(z, index))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": "tf.Tensor(\n[[0.  1.  0.  0.  0. ]\n [0.  0.  1.  0.  0. ]\n [0.  0.  0.  0.5 0.5]\n [1.  0.  0.  0.  0. ]\n [0.  0.  0.  0.  1. ]], shape=(5, 5), dtype=float32)\ntf.Tensor(\n[[0. ]\n [1. ]\n [2. ]\n [1.5]\n [6. ]], shape=(5, 1), dtype=float32)\ntf.Tensor([ 0.  0. -1.  0.  0.], shape=(5,), dtype=float32)\n"
    }
   ],
   "source": [
    "@tf.function\n",
    "def bulk_bandwidthify(indices, bandwidth):\n",
    "    num_indices = tf.shape(indices)[0]\n",
    "    \n",
    "    indices = tf.unstack(indices)\n",
    "    result = tf.zeros((num_indices, bandwidth), dtype=tf.float32)\n",
    "    result = tf.unstack(result)\n",
    "    \n",
    "    for i, index in enumerate(indices):\n",
    "        b_index = bandwidthify(index, bandwidth)\n",
    "        result[i] += b_index\n",
    "    \n",
    "    result = tf.stack(result)\n",
    "    return result\n",
    "\n",
    "indices = tf.Variable([1,2,3.5,0,4],dtype=tf.float32)\n",
    "bandwidth = tf.constant(5, dtype=tf.int32)\n",
    "dummy_array = tf.cast(tf.range(bandwidth), tf.float32)\n",
    "with tf.GradientTape() as tape:\n",
    "    z = bulk_bandwidthify(indices, bandwidth)\n",
    "    nz = superposition_lookup_vectored(dummy_array, z) # Lookup operation\n",
    "\n",
    "print(z)\n",
    "print(nz)\n",
    "print(tape.gradient(nz, indices))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Residual lookup\n",
    "In this method, we return two tensors, the result and the residue. So, although the result is not differentiable wrt to index, the residue is. This allows us to propagate some extra information in parallel which can then be consumed intelligently by some algorithm in downstream. This has the benefit that the result always exists in the original array and is never an interpolation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": "3 -0.5\n[0 0 1 0 0] [0 0 0 0 0]\n0 1\n"
    }
   ],
   "source": [
    "@tf.function\n",
    "def residual_lookup(arr, index):\n",
    "    i = tf.round(index)\n",
    "    residue = index - i\n",
    "    i = tf.cast(i, tf.int32)\n",
    "    \n",
    "    result = arr[i]\n",
    "    \n",
    "    return result, residue\n",
    "\n",
    "arr = tf.Variable([1,2,3,4,5],dtype=tf.float32)\n",
    "index = tf.Variable(1.5, dtype=tf.float32)\n",
    "\n",
    "with tf.GradientTape(persistent=True) as tape:\n",
    "    result, residue = residual_lookup(arr, index)\n",
    "\n",
    "tf.print(result, residue)\n",
    "tf.print(tape.gradient(result, arr), tape.gradient(residue, arr))\n",
    "tf.print(tape.gradient(result, index), tape.gradient(residue, index))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Asymmetrical Vectored Lookup\n",
    "\n",
    "In this method, we calculate the result of the forward pass by finding the most likely index of the vector and returning the value associated with that index. The forward pass is thus non-differentiable. Therefore we have to define our own backward pass. To calculate the backward pass, we first estimate our target. The $target$ is the difference of the result obtained in the forward pass and the gradients from the loss. We now, find the value in our vector which is closest to the target. We want to increase the probability of this index while decreasing the probability of other indexes. So, we create a vector which is $-1$ at the index of the target and $1$ everywhere else. The optimizer substracts the gradient, so it has to be negative.\n",
    "\n",
    "The gradients of the forward pass is not equal to the backward pass. In that sense, this is asymmetric."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": "[2 10]\n[[-1 1 1]\n [-1 1 1]]\n[[1 1 -1]\n [10 -10 10]]\n"
    }
   ],
   "source": [
    "@tf.function\n",
    "@tf.custom_gradient\n",
    "def asymmetrical_vectored_lookup(v, k):\n",
    "    k_shape = tf.shape(k)\n",
    "\n",
    "    # Pick the value at the most likely index, non-differentiably\n",
    "    b_idx = tf.argmax(k, axis=-1)\n",
    "    idx_len = tf.shape(b_idx)[0]\n",
    "    a_idx = tf.range(idx_len, dtype=tf.int64)\n",
    "    idx = tf.stack([a_idx, b_idx], axis=1)\n",
    "    forward_result = tf.gather_nd(v, idx)\n",
    "\n",
    "    def grad(upstream_grads):\n",
    "        # Estimate the target scalar which we want to look up\n",
    "        target = forward_result - upstream_grads\n",
    "        target = tf.expand_dims(target, -1)\n",
    "\n",
    "        # Find the index of element in the array which is closest to target\n",
    "        diff_vector = tf.math.squared_difference(v, target)\n",
    "        d_idx = tf.argmin(diff_vector, axis=-1)\n",
    "\n",
    "        # Create a vector which is 1 everywhere except the idx\n",
    "        # of the target, where it is -1\n",
    "        ones = tf.ones(k_shape)\n",
    "        eyes = tf.one_hot([d_idx], k_shape[-1])[0]\n",
    "        k_grad = -(2 * eyes - ones)\n",
    "\n",
    "        # d/dv (v . k) = k\n",
    "        v_grad = k\n",
    "\n",
    "        upstream_grads = tf.expand_dims(upstream_grads, -1)\n",
    "        return upstream_grads * v_grad, tf.math.abs(upstream_grads) * k_grad\n",
    "\n",
    "    return forward_result, grad\n",
    "\n",
    "v = tf.constant([[1,2,3], [10,20,30]], dtype=tf.float32)\n",
    "k = tf.constant([[0,1,0], [1,0,0]], dtype=tf.float32)\n",
    "t = tf.constant([3, 20], dtype=tf.float32)\n",
    "\n",
    "with tf.GradientTape(persistent=True) as tape:\n",
    "    tape.watch(k)\n",
    "    r = asymmetrical_vectored_lookup(v, k)\n",
    "    loss = tf.nn.l2_loss(r - t)\n",
    "\n",
    "tf.print(r)\n",
    "tf.print(tape.gradient(r, k))\n",
    "tf.print(tape.gradient(loss, k))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": "2000 [-0 -0] [20 10]\n2000 [0.577248037 0.577248216] [20 10]\n2000 [0.542812288 0.54283917] [20 10]\n2000 [0.822554886 0.822548866] [20 10]\n2000 [0.728512 0.728538215] [20 10]\n0 [0.875905 0.875911474] [40 70]\n0 [0.686750948 0.686771] [40 70]\n0 [0.781492531 0.781499624] [40 70]\n0 [0.552687764 0.552703917] [40 70]\n0 [0.624352634 0.624367058] [40 70]\n"
    },
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "<tf.Tensor: shape=(2, 3), dtype=float32, numpy=\narray([[12., 88.,  0.],\n       [ 0., 88., 12.]], dtype=float32)>"
     },
     "metadata": {},
     "execution_count": 12
    }
   ],
   "source": [
    "values = tf.constant([\n",
    "    [20,40,30],\n",
    "    [50,70,10],\n",
    "], dtype=tf.float32)\n",
    "choice = tf.Variable([\n",
    "    [1,0,0],\n",
    "    [0,0,1],\n",
    "], dtype=tf.float32)\n",
    "\n",
    "target = tf.constant([40, 70], dtype=tf.float32)\n",
    "\n",
    "# opt = tf.keras.optimizers.Adam(3e-4)\n",
    "# opt = tf.keras.optimizers.Adam(1e-2)\n",
    "lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(\n",
    "    1e-1,\n",
    "    decay_steps=100,\n",
    "    decay_rate=1e-1,\n",
    "    staircase=True)\n",
    "opt = tf.keras.optimizers.Adam(lr_schedule)\n",
    "\n",
    "steps = 10\n",
    "\n",
    "for i in range(steps):\n",
    "    with tf.GradientTape() as tape:\n",
    "        out = asymmetrical_vectored_lookup(values, choice)\n",
    "        target_loss = tf.nn.l2_loss(out - target)\n",
    "        entropy_loss = entropy(choice)\n",
    "\n",
    "        loss = target_loss + entropy_loss * 1e-2\n",
    "\n",
    "    variables = [choice]\n",
    "    grads = tape.gradient(loss, variables)\n",
    "\n",
    "    opt.apply_gradients(zip(grads, variables))\n",
    "    choice.assign(to_prob_dist_all(choice))\n",
    "\n",
    "    if i % (steps // 10) == 0:\n",
    "        tf.print(target_loss, entropy_loss, out, )\n",
    "\n",
    "tf.round(choice * 100)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Array assignment\n",
    "Tensorflow does not support direct index assignment of variables. So, instead we use a masking technique"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": "tf.Tensor(\n[[1. 1. 1.]\n [4. 4. 4.]\n [3. 3. 3.]], shape=(3, 3), dtype=float32)\ntf.Tensor(\n[[1. 1. 1.]\n [0. 0. 0.]\n [1. 1. 1.]], shape=(3, 3), dtype=float32)\nWARNING:tensorflow:The dtype of the source tensor must be floating (e.g. tf.float32) when calling GradientTape.gradient, got tf.int32\nNone\ntf.Tensor([1. 1. 1.], shape=(3,), dtype=float32)\n"
    }
   ],
   "source": [
    "@tf.function\n",
    "def assign_index(arr, index, element):\n",
    "    arr_shape = tf.shape(arr)\n",
    "    \n",
    "    pos_mask = tf.eye(arr_shape[0])[index]\n",
    "    pos_mask = tf.transpose(tf.expand_dims(pos_mask, 0))\n",
    "    neg_mask = 1 - pos_mask\n",
    "    \n",
    "    tiled_element = tf.reshape(tf.tile(element, [arr_shape[0]]), arr_shape)\n",
    "    \n",
    "    arr = arr * neg_mask + tiled_element * pos_mask\n",
    "    \n",
    "    return arr\n",
    "\n",
    "arr = tf.Variable([\n",
    "    [1,1,1],\n",
    "    [2,2,2],\n",
    "    [3,3,3]\n",
    "],dtype=tf.float32)\n",
    "index = tf.constant(1)\n",
    "element = tf.Variable([4,4,4],dtype=tf.float32)\n",
    "\n",
    "with tf.GradientTape(persistent=True) as tape:\n",
    "    new_arr = assign_index(arr, index, element)\n",
    "    \n",
    "print(new_arr)\n",
    "print(tape.gradient(new_arr, arr))\n",
    "print(tape.gradient(new_arr, index))\n",
    "print(tape.gradient(new_arr, element))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Superpositioned assignment in case of a vector like index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": "[[1 1 1]\n [4 4 4]\n [3 3 3]] [9 6 3]\n[[2.5 2.5 2.5]\n [3 3 3]\n [3 3 3]] [9 6 3]\ntf.Tensor(\n[[1. 1. 1.]\n [0. 0. 0.]\n [1. 1. 1.]], shape=(3, 3), dtype=float32)\ntf.Tensor([1. 1. 1.], shape=(3,), dtype=float32)\n"
    }
   ],
   "source": [
    "@tf.function\n",
    "def assign_index_vectored(arr, index, element):\n",
    "    arr_shape = tf.shape(arr)\n",
    "    \n",
    "    pos_mask = tf.transpose(tf.expand_dims(index, 0))\n",
    "    neg_mask = 1 - pos_mask\n",
    "    \n",
    "    tiled_element = tf.reshape(tf.tile(element, [arr_shape[0]]), arr_shape)\n",
    "\n",
    "    arr = arr * neg_mask + tiled_element * pos_mask\n",
    "    \n",
    "    return arr\n",
    "\n",
    "arr = tf.Variable([\n",
    "    [1,1,1],\n",
    "    [2,2,2],\n",
    "    [3,3,3]\n",
    "],dtype=tf.float32)\n",
    "index1 = tf.Variable([0,1,0], dtype=tf.float32)\n",
    "index2 = tf.Variable([0.5,0.5,0], dtype=tf.float32)\n",
    "element = tf.Variable([4,4,4],dtype=tf.float32)\n",
    "\n",
    "with tf.GradientTape(persistent=True) as tape:\n",
    "    new_arr1 = assign_index_vectored(arr, index1, element)\n",
    "    new_arr2 = assign_index_vectored(arr, index2, element)\n",
    "\n",
    "tf.print(new_arr1, tape.gradient(new_arr1, index1))\n",
    "tf.print(new_arr2, tape.gradient(new_arr2, index2))\n",
    "print(tape.gradient(new_arr1, arr))\n",
    "print(tape.gradient(new_arr1, element))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Higher dimensional arrays"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "<tf.Tensor: shape=(3, 3, 2), dtype=float32, numpy=\narray([[[6., 6.],\n        [6., 6.],\n        [6., 6.]],\n\n       [[6., 6.],\n        [6., 6.],\n        [6., 6.]],\n\n       [[6., 6.],\n        [6., 6.],\n        [6., 6.]]], dtype=float32)>"
     },
     "metadata": {},
     "execution_count": 15
    }
   ],
   "source": [
    "@tf.function\n",
    "def match_shapes(x, y):\n",
    "    # Find which one needs to be broadcasted\n",
    "    low, high = (y, x) if tf.rank(x) > tf.rank(y) else (x, y)\n",
    "    l_rank, l_shape = tf.rank(low), tf.shape(low)\n",
    "    h_rank, h_shape = tf.rank(high), tf.shape(high)\n",
    "    \n",
    "    # Find the difference in ranks\n",
    "    common_shape = h_shape[:l_rank]\n",
    "    tf.debugging.assert_equal(common_shape, l_shape, 'No common shape to broadcast')\n",
    "    padding = tf.ones(h_rank - l_rank, dtype=tf.int32)\n",
    "    \n",
    "    # Pad the difference with ones and reshape\n",
    "    new_shape = tf.concat((common_shape, padding),axis=0)\n",
    "    low = tf.reshape(low, new_shape)\n",
    "\n",
    "    return high, low\n",
    "\n",
    "@tf.function\n",
    "def broadcast_multiply(x, y):\n",
    "    x, y = match_shapes(x, y)\n",
    "    return x * y\n",
    "    \n",
    "x = tf.ones((3, 3, 2)) * 3\n",
    "y = tf.ones((3, 3)) * 2\n",
    "broadcast_multiply(x, y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": "tf.Tensor([  2. 222.], shape=(2,), dtype=float32)\ntf.Tensor(\n[[[0. 0.]\n  [0. 0.]\n  [0. 0.]]\n\n [[0. 0.]\n  [0. 0.]\n  [1. 1.]]\n\n [[0. 0.]\n  [0. 0.]\n  [0. 0.]]], shape=(3, 3, 2), dtype=float32)\ntf.Tensor([  0. 224.   0.], shape=(3,), dtype=float32)\ntf.Tensor([  0.   0. 224.], shape=(3,), dtype=float32)\n"
    }
   ],
   "source": [
    "@tf.function\n",
    "def tensor_lookup_2d(arr, x_index, y_index):\n",
    "    # Calculate outer product\n",
    "    mask = tf.tensordot(x_index, y_index, axes=0)\n",
    "    \n",
    "    # Broadcast the mask to match dimensions with arr\n",
    "    masked_arr = broadcast_multiply(mask, arr)\n",
    "    \n",
    "    # Reduce max to extract the cell\n",
    "    element = tf.math.reduce_max(masked_arr, axis=[0,1])\n",
    "    return element\n",
    "\n",
    "arr = tf.Variable([\n",
    "    [[1,1],[1,11],[1,111]],\n",
    "    [[2,2],[2,22],[2,222]],\n",
    "    [[3,3],[3,33],[3,333]]\n",
    "],dtype=tf.float32)\n",
    "x_index = tf.Variable(tf.one_hot(1, 3),dtype=tf.float32)\n",
    "y_index = tf.Variable(tf.one_hot(2, 3),dtype=tf.float32)\n",
    "\n",
    "with tf.GradientTape(persistent=True) as tape:\n",
    "    element = tensor_lookup_2d(arr, x_index, y_index)\n",
    "    \n",
    "print(element)\n",
    "print(tape.gradient(element, arr))\n",
    "print(tape.gradient(element, x_index))\n",
    "print(tape.gradient(element, y_index))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": "tf.Tensor(\n[[[  1.   1.]\n  [  1.  11.]\n  [  1. 111.]]\n\n [[  2.   2.]\n  [  2.  22.]\n  [  5. 555.]]\n\n [[  3.   3.]\n  [  3.  33.]\n  [  3. 333.]]], shape=(3, 3, 2), dtype=float32)\ntf.Tensor(\n[[[1. 1.]\n  [1. 1.]\n  [1. 1.]]\n\n [[1. 1.]\n  [1. 1.]\n  [0. 0.]]\n\n [[1. 1.]\n  [1. 1.]\n  [1. 1.]]], shape=(3, 3, 2), dtype=float32)\ntf.Tensor([1. 1.], shape=(2,), dtype=float32)\ntf.Tensor([448. 336. 224.], shape=(3,), dtype=float32)\ntf.Tensor([556. 536. 336.], shape=(3,), dtype=float32)\n"
    }
   ],
   "source": [
    "@tf.function\n",
    "def tensor_write_2d(arr, element, x_index, y_index):\n",
    "    arr_shape = tf.shape(arr)\n",
    "    mask = tf.tensordot(x_index, y_index, axes=0)\n",
    "    \n",
    "    # Broadcast the mask to match dimensions with arr\n",
    "    _, mask = match_shapes(arr, mask)\n",
    "    \n",
    "    element = tf.reshape(element,[1,1,-1])\n",
    "    element = tf.tile(element, [arr_shape[0], arr_shape[1], 1])\n",
    "    \n",
    "    result = (1.0 - mask) * arr + mask * element\n",
    "    \n",
    "    return result\n",
    "\n",
    "arr = tf.Variable([\n",
    "    [[1,1],[1,11],[1,111]],\n",
    "    [[2,2],[2,22],[2,222]],\n",
    "    [[3,3],[3,33],[3,333]]\n",
    "],dtype=tf.float32)\n",
    "element = tf.Variable([5,555], dtype=tf.float32)\n",
    "x_index = tf.Variable(tf.one_hot(1, 3),dtype=tf.float32)\n",
    "y_index = tf.Variable(tf.one_hot(2, 3),dtype=tf.float32)\n",
    "\n",
    "with tf.GradientTape(persistent=True) as tape:\n",
    "    new_arr = tensor_write_2d(arr, element, x_index, y_index)\n",
    "    \n",
    "print(new_arr)\n",
    "print(tape.gradient(new_arr, arr))\n",
    "print(tape.gradient(new_arr, element))\n",
    "print(tape.gradient(new_arr, x_index))\n",
    "print(tape.gradient(new_arr, y_index))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.6-final"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}