{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "import random\n",
    "from math import pi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['🟥', '🟦', '🟧', '🟨', '🟩', '🟪', '🟫']\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(7, 2), dtype=float32, numpy=\n",
       "array([[ 0.9848077 ,  0.17364818],\n",
       "       [-0.777146  , -0.62932044],\n",
       "       [ 0.7880108 ,  0.6156615 ],\n",
       "       [ 0.6691306 ,  0.74314487],\n",
       "       [ 0.17364822,  0.9848077 ],\n",
       "       [ 0.20791148, -0.9781476 ],\n",
       "       [ 0.9396926 ,  0.34202012]], dtype=float32)>"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "squares = [chr(i) for i in range(0x1F7E5, 0x1F7EC)]\n",
    "tfsquares = tf.constant(squares)\n",
    "print(squares)\n",
    "colors = tf.constant([10, 219, 38, 48, 80,  282, 20] * tf.constant(pi/180.0), dtype=tf.float32)\n",
    "color_vectors = tf.transpose(\n",
    "    tf.stack([\n",
    "    tf.math.cos(colors),\n",
    "    tf.math.sin(colors)]\n",
    "    )\n",
    ")\n",
    "color_vectors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "def one_hot_to_string(tensor):\n",
    "    square_select = tf.math.argmax(tensor, 2)\n",
    "\n",
    "    tfstring = tf.strings.join(\n",
    "        tf.map_fn(\n",
    "            lambda v:\n",
    "            tf.strings.join(\n",
    "                tf.gather(tfsquares, tf.cast(v, tf.int64))), square_select, fn_output_signature=tf.string), \"\\n\")\n",
    "\n",
    "    return tfstring.numpy().decode()\n",
    "\n",
    "def index_to_string(tensor):\n",
    "    return tf.argmax((tf.einsum(\"ijk,lk->ijl\", tensor, color_vectors)), axis=2)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(8, 8), dtype=int64, numpy=\n",
       "array([[5, 0, 1, 3, 3, 2, 4, 0],\n",
       "       [3, 2, 3, 6, 5, 3, 0, 5],\n",
       "       [1, 3, 0, 1, 6, 4, 5, 3],\n",
       "       [3, 5, 0, 6, 5, 5, 0, 3],\n",
       "       [0, 3, 0, 6, 6, 1, 5, 2],\n",
       "       [4, 1, 4, 6, 6, 1, 0, 5],\n",
       "       [6, 0, 2, 5, 0, 1, 5, 4],\n",
       "       [0, 2, 5, 4, 1, 5, 0, 4]])>"
      ]
     },
     "execution_count": 57,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vi = tf.map_fn(lambda v: tf.map_fn(lambda s: color_vectors[s],v, fn_output_signature=tf.float32),tf.constant(e, shape=(8, 8)),fn_output_signature=tf.float32)\n",
    "\n",
    "index_to_string(vi)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(1, 8, 8), dtype=float32, numpy=\n",
       "array([[[0.71719116, 0.04602498, 1.0273366 , 2.9402392 , 2.4437628 ,\n",
       "         1.0988994 , 1.7921778 , 0.6444071 ],\n",
       "        [1.2672982 , 1.8689629 , 3.8761568 , 2.4126112 , 2.840665  ,\n",
       "         3.0419385 , 4.8909955 , 3.583321  ],\n",
       "        [0.4506898 , 0.01335579, 3.258954  , 3.80125   , 6.4259176 ,\n",
       "         6.3595123 , 6.609476  , 1.6973627 ],\n",
       "        [1.8337119 , 1.4464988 , 4.6484404 , 4.2732587 , 7.6358337 ,\n",
       "         5.4232316 , 4.313001  , 1.6071305 ],\n",
       "        [4.9838686 , 1.3140092 , 1.5253077 , 1.9003649 , 6.782009  ,\n",
       "         3.5835938 , 4.203249  , 1.5495552 ],\n",
       "        [5.3104043 , 5.0123706 , 3.3745103 , 4.540594  , 6.6606255 ,\n",
       "         3.9076035 , 4.9301    , 3.4342353 ],\n",
       "        [3.3523657 , 4.9394774 , 5.672796  , 4.8756404 , 7.1605525 ,\n",
       "         8.192585  , 7.06907   , 4.199152  ],\n",
       "        [0.13434029, 4.039145  , 6.2488484 , 3.2168784 , 5.319997  ,\n",
       "         5.4898353 , 5.301211  , 1.6475667 ]]], dtype=float32)>"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# def colors_to_one_hot(tensor):\n",
    "vi = tf.map_fn(lambda v: tf.map_fn(lambda s: color_vectors[s],v, fn_output_signature=tf.float32),tf.constant(e, shape=(8, 8)),fn_output_signature=tf.float32)\n",
    "ns = tf.image.extract_patches(tf.constant(vi, shape=(1,8,8,2)), (1,3,3,1), (1,1,1,1), (1,1,1,1), padding=\"SAME\")\n",
    "virs = tf.reshape(vi, (8*8,2))\n",
    "# tf.reduce_sum(tf.multiply(tf.repeat(vi, 9, -1), ns))\n",
    "tf.abs(tf.einsum(\"...k,...k->...\", tf.repeat(vi, 9, -1), ns))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tf.Tensor([0. 0. 0. 1. 0. 0. 0.], shape=(7,), dtype=float32)\n",
      "tf.Tensor([[3]], shape=(1, 1), dtype=int64)\n",
      "🟨🟥🟨🟨🟩🟫🟪🟫\n",
      "🟩🟪🟧🟧🟦🟧🟦🟫\n",
      "🟩🟦🟫🟨🟦🟨🟩🟫\n",
      "🟪🟫🟩🟨🟧🟪🟥🟩\n",
      "🟥🟫🟨🟫🟧🟦🟦🟧\n",
      "🟩🟫🟫🟧🟨🟨🟪🟫\n",
      "🟩🟪🟦🟥🟧🟫🟧🟨\n",
      "🟨🟪🟨🟧🟪🟦🟨🟧\n",
      "--------------\n",
      "🟧🟨🟦🟪🟧🟨🟪🟨\n",
      "🟨🟧🟫🟧🟥🟦🟪🟩\n",
      "🟫🟪🟨🟨🟧🟫🟫🟩\n",
      "🟧🟦🟦🟧🟫🟨🟫🟥\n",
      "🟩🟥🟪🟧🟨🟩🟫🟪\n",
      "🟫🟩🟨🟦🟨🟫🟦🟩\n",
      "🟫🟦🟧🟦🟧🟧🟪🟩\n",
      "🟫🟪🟫🟩🟨🟨🟥🟨\n",
      "tf.Tensor(\n",
      "[[0. 0. 0. 0. 0. 0. 0. 0.]\n",
      " [0. 0. 1. 0. 0. 1. 0. 0.]\n",
      " [0. 0. 0. 0. 0. 0. 0. 0.]\n",
      " [0. 0. 0. 0. 0. 0. 0. 0.]\n",
      " [0. 0. 0. 0. 0. 0. 0. 0.]\n",
      " [0. 0. 0. 0. 0. 0. 0. 0.]\n",
      " [0. 0. 0. 0. 0. 0. 0. 0.]\n",
      " [0. 0. 0. 0. 0. 0. 0. 0.]], shape=(8, 8), dtype=float32)\n"
     ]
    }
   ],
   "source": [
    "oh = tf.one_hot(3, 7)\n",
    "print(oh)\n",
    "print(tf.where(oh))\n",
    "\n",
    "# a = tf.constant([[list(tf.one_hot(random.randrange(7), 7)) for i in range(8) ] for i in range(8)], shape=(8,8,len(squares)))\n",
    "\n",
    "# print([[list(tf.one_hot(random.randrange(7), 7)) for i in range(8) ] for i in range(8)])\n",
    "e = [[random.randrange(7) for i in range(8)] for i in range(8)]\n",
    "# print(e)\n",
    "\n",
    "a = tf.one_hot(tf.constant(e, shape=(8, 8)), len(squares))\n",
    "\n",
    "print(one_hot_to_string(a))\n",
    "# print(\"--------------\")\n",
    "# print(one_hot_to_string(tf.reverse(a, axis=(-2,))))\n",
    "# print(\"--------------\")\n",
    "# print(one_hot_to_string(tf.reverse(a, axis=(-3,))))\n",
    "# print(\"--------------\")\n",
    "# print(one_hot_to_string(tf.transpose(a, perm=(1,0,2))))\n",
    "print(\"--------------\")\n",
    "# print(one_hot_to_string(tf.reverse(tf.transpose(tf.reverse(a, axis=(-2,)), perm=(1,0,2)), axis=(-2,))))\n",
    "print(one_hot_to_string(tf.reverse(a, axis=(-2,-3,))))\n",
    "# print(tensor_to_string(tf.transpose(a, [1, 0, 2])))\n",
    "# print(\"--------------\")\n",
    "# print(tensor_to_string(tf.linalg.matmul( a , tf.constant(\n",
    "#     [\n",
    "#         [0, 0, 0, 0, 0, 0, 0, 1],\n",
    "#         [0, 0, 0, 0, 0, 0, 1, 0],\n",
    "#         [0, 0, 0, 0, 0, 1, 0, 0],\n",
    "#         [0, 0, 0, 0, 1, 0, 0, 0],\n",
    "#         [0, 0, 0, 1, 0, 0, 0, 0],\n",
    "#         [0, 0, 1, 0, 0, 0, 0, 0],\n",
    "#         [0, 1, 0, 0, 0, 0, 0, 0],\n",
    "#         [1, 0, 0, 0, 0, 0, 0, 0]\n",
    "#     ], dtype=tf.float32, shape=(8,8,1)\n",
    "# ), transpose_a=True)))\n",
    "# print(\"--------------\")\n",
    "# print(tensor_to_string(tf.reverse(a, (0,))))\n",
    "# print(\"--------------\")\n",
    "# print(tensor_to_string(tf.reverse(a, (1,))))\n",
    "# print(tensor_to_string(\n",
    "#     tf.linalg.matmul(\n",
    "#         tf.transpose(a, [0, 2, 1]),\n",
    "#         tf.reverse(a, (1,)))\n",
    "# ))\n",
    "print(\n",
    "    tf.einsum(\n",
    "        \"ijk,ijk->ij\",\n",
    "        a,\n",
    "        tf.reverse(a, (1,)))\n",
    ")\n",
    "\n",
    "\n",
    "# print(tf.constant(range(6), shape=(2,3)))\n",
    "# tf.map_fn(tf.math.reduce_sum, tf.constant(range(6), shape=(2,3)))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(), dtype=float32, numpy=-0.08645747>"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import aesthetic_loss\n",
    "\n",
    "es = [[[random.randrange(7) for i in range(8)]\n",
    "      for i in range(8)] for j in range(16)]\n",
    "\n",
    "vis = tf.stack(\n",
    "    [tf.map_fn(lambda v: tf.map_fn(lambda s: color_vectors[s], v, fn_output_signature=tf.float32), tf.constant(\n",
    "        e, shape=(8, 8)), fn_output_signature=tf.float32) for e in es]\n",
    ")\n",
    "\n",
    "aesthetic_loss.aesthetic_loss(vis)\n",
    "# aesthetic_loss.compute_score(vis, vis)\n",
    "# tf.linalg.l2_normalize(vis, axis=-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 91,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(), dtype=int32, numpy=2890>"
      ]
     },
     "execution_count": 91,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x = tf.Variable(2)\n",
    "y = tf.Variable(3)\n",
    "\n",
    "xp = tf.pow(x, [0,1,2,3])\n",
    "yp = tf.pow(y, [0,1,2,3])\n",
    "\n",
    "A = tf.Variable([[i+j for i in range(4)] for j in range(4)], shape=(4,4))\n",
    "tf.einsum(\"i,ij,j\", xp, A, yp)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 70,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(16, 8), dtype=float32, numpy=\n",
       "array([[-6.07829041e+02, -1.43090210e+03, -2.25397559e+03,\n",
       "        -3.07704883e+03, -3.90012207e+03, -4.72319531e+03,\n",
       "        -5.54626855e+03, -6.36934229e+03],\n",
       "       [-4.28289948e+01, -8.77924423e+01, -1.32755875e+02,\n",
       "        -1.77719330e+02, -2.22682800e+02, -2.67646210e+02,\n",
       "        -3.12609680e+02, -3.57573120e+02],\n",
       "       [ 9.78072433e+01,  2.85662659e+02,  4.73518066e+02,\n",
       "         6.61373413e+02,  8.49228882e+02,  1.03708435e+03,\n",
       "         1.22493982e+03,  1.41279517e+03],\n",
       "       [ 3.64326048e+00,  2.58858261e+01,  4.81283913e+01,\n",
       "         7.03709641e+01,  9.26135254e+01,  1.14856094e+02,\n",
       "         1.37098648e+02,  1.59341217e+02],\n",
       "       [-1.08920467e+00,  3.87969995e+00,  8.84860516e+00,\n",
       "         1.38175116e+01,  1.87864113e+01,  2.37553177e+01,\n",
       "         2.87242279e+01,  3.36931267e+01],\n",
       "       [ 2.75129032e+01,  9.32729721e+01,  1.59033051e+02,\n",
       "         2.24793137e+02,  2.90553192e+02,  3.56313324e+02,\n",
       "         4.22073334e+02,  4.87833435e+02],\n",
       "       [-5.32951474e-01,  8.84730148e+00,  1.82275543e+01,\n",
       "         2.76078072e+01,  3.69880600e+01,  4.63683128e+01,\n",
       "         5.57485619e+01,  6.51288223e+01],\n",
       "       [-7.10889587e+01, -1.97856873e+02, -3.24624786e+02,\n",
       "        -4.51392792e+02, -5.78160645e+02, -7.04928589e+02,\n",
       "        -8.31696533e+02, -9.58464539e+02],\n",
       "       [-3.13533902e-01,  1.77808132e+01,  3.58751602e+01,\n",
       "         5.39695053e+01,  7.20638504e+01,  9.01582031e+01,\n",
       "         1.08252541e+02,  1.26346893e+02],\n",
       "       [ 7.97924957e+01,  2.76383240e+02,  4.72974030e+02,\n",
       "         6.69564819e+02,  8.66155518e+02,  1.06274634e+03,\n",
       "         1.25933704e+03,  1.45592786e+03],\n",
       "       [-2.12144566e+01, -1.41624451e+02, -2.62034424e+02,\n",
       "        -3.82444489e+02, -5.02854431e+02, -6.23264404e+02,\n",
       "        -7.43674438e+02, -8.64084229e+02],\n",
       "       [-1.87502289e+02, -3.70114594e+02, -5.52726929e+02,\n",
       "        -7.35339233e+02, -9.17951660e+02, -1.10056396e+03,\n",
       "        -1.28317627e+03, -1.46578857e+03],\n",
       "       [-1.07772827e+00, -9.14207935e-01, -7.50688553e-01,\n",
       "        -5.87172508e-01, -4.23656583e-01, -2.60130525e-01,\n",
       "        -9.66115594e-02,  6.69060946e-02],\n",
       "       [ 2.14334583e+01,  1.10324554e+02,  1.99215652e+02,\n",
       "         2.88106781e+02,  3.76997894e+02,  4.65889008e+02,\n",
       "         5.54780090e+02,  6.43671143e+02],\n",
       "       [ 5.15513878e+01,  1.88147568e+02,  3.24743774e+02,\n",
       "         4.61339966e+02,  5.97936157e+02,  7.34532349e+02,\n",
       "         8.71128601e+02,  1.00772473e+03],\n",
       "       [ 2.68221771e+02,  7.04315491e+02,  1.14040918e+03,\n",
       "         1.57650305e+03,  2.01259656e+03,  2.44869043e+03,\n",
       "         2.88478394e+03,  3.32087769e+03]], dtype=float32)>"
      ]
     },
     "execution_count": 70,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "r = tf.random.normal((16,2))\n",
    "rp = tf.map_fn(lambda v: tf.pow(tf.repeat(tf.reshape(v, shape=(2,1)), repeats=4, axis=1), tf.constant([[0,1,2,3], [0,1,2,3]], dtype=tf.float32)), r)\n",
    "A = tf.Variable(tf.reshape([float(i) for i in range(8*4*4)], shape=(8,4,4)), shape=(8,4,4), dtype=tf.float32)\n",
    "tf.einsum(\"ik,ij,lkj->il\", rp[:,0], rp[:,1], A)\n",
    "# rp[:,0]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 82,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tf.Tensor(\n",
      "[[ 4.15302128e-01 -1.17927468e+00  2.49658436e-01  6.19551420e-01\n",
      "  -1.39033377e-01  2.79209971e-01  3.97601783e-01 -2.17277408e-01]\n",
      " [ 1.74276757e+00  5.69359064e-02  5.38259745e-04  1.49045396e+00\n",
      "  -7.87080288e-01 -3.98119450e-01  1.18716180e-01  1.86502957e+00]\n",
      " [ 1.78081155e+00 -1.29199183e+00 -5.72377980e-01  1.73489356e+00\n",
      "  -1.81523621e+00 -7.56731749e-01  5.50791979e-01  2.23066378e+00]\n",
      " [-2.86107287e-02  2.38285348e-01 -2.90394664e-01  2.20198780e-02\n",
      "   3.43291536e-02  2.77290910e-01 -3.25144343e-02  3.50418538e-01]\n",
      " [-2.32135087e-01  1.95894271e-01  5.75335659e-02 -1.76543556e-03\n",
      "  -4.01449911e-02  2.00845733e-01 -9.64274257e-02  2.31855541e-01]\n",
      " [ 4.60646629e-01  7.75275290e-01 -9.96064782e-01 -5.15016496e-01\n",
      "   2.03099363e-02  7.87942708e-02 -3.07008505e-01  7.16598153e-01]\n",
      " [ 3.02803397e-01 -8.83126408e-02  3.72119397e-01  1.14753373e-01\n",
      "  -8.30153376e-03  6.69163764e-02 -3.60930443e-01  4.32033747e-01]\n",
      " [-1.55036688e-01  2.96032988e-02  5.69820963e-02 -1.04342289e-01\n",
      "   2.18345195e-01  1.60513148e-01 -2.30654161e-02  2.92460948e-01]\n",
      " [-9.37448144e-02  4.92064208e-02  1.25066668e-01 -1.08998209e-01\n",
      "   1.47912696e-01  7.13704750e-02 -9.69988406e-02  3.16246897e-01]\n",
      " [ 3.76237392e-01 -8.38992000e-02  3.31230164e-01  2.86825836e-01\n",
      "   3.67110878e-01 -3.61552715e-01  6.06793404e-01 -1.22976625e+00]\n",
      " [-1.00743435e-01 -1.46258920e-01 -1.82411104e-01 -3.14753801e-01\n",
      "   7.92229354e-01  2.98170269e-01 -1.82591528e-01  2.78764069e-01]\n",
      " [ 3.58289123e-01  1.97965109e+00 -1.18626821e+00 -1.47391105e+00\n",
      "   6.11245155e-01 -4.61193591e-01 -7.64733970e-01  1.45492101e+00]\n",
      " [-2.00078160e-01  4.32360053e-01  1.63023725e-01 -3.31694707e-02\n",
      "  -4.53773856e-01  4.14216101e-01  1.76714227e-01  2.12158859e-01]\n",
      " [ 4.91659880e-01  1.56133473e+00 -8.19299459e-01  9.75016415e-01\n",
      "  -2.52336192e+00  1.69403172e+00  2.05944443e+00  1.03490460e+00]\n",
      " [-2.33240795e+00  6.64168453e+00 -4.04779243e+00 -5.99489832e+00\n",
      "   1.60179734e+00  2.19676828e+00  4.11457443e+00  2.81916380e-01]\n",
      " [ 1.73429586e-02  4.88152355e-02  3.41429591e-01 -2.39478126e-02\n",
      "  -2.88104061e-02  1.31737068e-01 -2.55863100e-01  2.94781595e-01]], shape=(16, 8), dtype=float32)\n"
     ]
    }
   ],
   "source": [
    "import polymap\n",
    "r = tf.random.normal((16,2))\n",
    "\n",
    "pm = polymap.Polymap(8,4)\n",
    "# pm.build([16,2])\n",
    "print(pm(r))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 111,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(16, 8, 8, 2), dtype=float32, numpy=\n",
       "array([[[[ 0.26760346,  0.9635291 ],\n",
       "         [ 0.92846906,  0.37140968],\n",
       "         [ 0.43006882,  0.90279603],\n",
       "         ...,\n",
       "         [ 0.9174245 ,  0.39790976],\n",
       "         [ 0.68369484,  0.729768  ],\n",
       "         [ 0.5905379 , -0.8070098 ]],\n",
       "\n",
       "        [[-0.9665002 ,  0.25666556],\n",
       "         [ 0.6174854 , -0.7865824 ],\n",
       "         [-0.8356619 , -0.54924417],\n",
       "         ...,\n",
       "         [-0.42570287,  0.90486294],\n",
       "         [ 0.99368197,  0.11223278],\n",
       "         [ 0.22659355,  0.97398937]],\n",
       "\n",
       "        [[ 0.9998629 ,  0.01655844],\n",
       "         [-0.02171519,  0.9997642 ],\n",
       "         [ 0.765471  , -0.64347047],\n",
       "         ...,\n",
       "         [-0.45096368,  0.8925422 ],\n",
       "         [ 0.23687297,  0.9715406 ],\n",
       "         [-0.2564354 ,  0.9665613 ]],\n",
       "\n",
       "        ...,\n",
       "\n",
       "        [[ 0.99743736, -0.07154454],\n",
       "         [-0.421027  ,  0.907048  ],\n",
       "         [-0.06904475,  0.99761355],\n",
       "         ...,\n",
       "         [-0.14730097, -0.9890917 ],\n",
       "         [-0.76936895, -0.63880455],\n",
       "         [ 0.12321126,  0.99238044]],\n",
       "\n",
       "        [[ 0.8592728 ,  0.51151776],\n",
       "         [ 0.5945256 , -0.80407673],\n",
       "         [-0.926837  , -0.37546387],\n",
       "         ...,\n",
       "         [-0.6461148 ,  0.76324016],\n",
       "         [-0.30289924,  0.95302254],\n",
       "         [-0.17756721,  0.9841086 ]],\n",
       "\n",
       "        [[-0.9989693 ,  0.04539123],\n",
       "         [ 0.9668924 ,  0.25518435],\n",
       "         [-0.70309734,  0.7110935 ],\n",
       "         ...,\n",
       "         [ 0.9820237 , -0.1887575 ],\n",
       "         [-0.45441642, -0.8907894 ],\n",
       "         [ 0.20170917,  0.9794454 ]]],\n",
       "\n",
       "\n",
       "       [[[-0.9041683 , -0.4271764 ],\n",
       "         [-0.8145172 ,  0.5801394 ],\n",
       "         [-0.93892664,  0.34411758],\n",
       "         ...,\n",
       "         [ 0.7724613 , -0.63506186],\n",
       "         [ 0.73973364,  0.6728997 ],\n",
       "         [-0.7584759 , -0.6517011 ]],\n",
       "\n",
       "        [[ 0.40460777,  0.91449016],\n",
       "         [ 0.9511079 ,  0.30885857],\n",
       "         [-0.8532944 ,  0.5214295 ],\n",
       "         ...,\n",
       "         [ 0.3517043 , -0.9361111 ],\n",
       "         [-0.81593966, -0.57813704],\n",
       "         [-0.96846116, -0.24916425]],\n",
       "\n",
       "        [[ 0.9971029 , -0.07606225],\n",
       "         [ 0.8346386 , -0.550798  ],\n",
       "         [ 0.79441607,  0.60737395],\n",
       "         ...,\n",
       "         [-0.17288695, -0.9849416 ],\n",
       "         [ 0.34486932, -0.93865067],\n",
       "         [ 0.15837449,  0.9873792 ]],\n",
       "\n",
       "        ...,\n",
       "\n",
       "        [[-0.9557019 ,  0.2943367 ],\n",
       "         [-0.1128222 , -0.9936151 ],\n",
       "         [ 0.8779246 , -0.4787989 ],\n",
       "         ...,\n",
       "         [ 0.01654347,  0.9998631 ],\n",
       "         [-0.31221828,  0.9500104 ],\n",
       "         [-0.70988506,  0.70431745]],\n",
       "\n",
       "        [[ 0.59169656, -0.8061608 ],\n",
       "         [ 0.29944003,  0.95411515],\n",
       "         [ 0.12906635,  0.9916359 ],\n",
       "         ...,\n",
       "         [-0.7345422 , -0.678563  ],\n",
       "         [-0.6101673 , -0.79227257],\n",
       "         [ 0.63129115,  0.7755459 ]],\n",
       "\n",
       "        [[ 0.15320726, -0.9881941 ],\n",
       "         [ 0.8636006 , -0.5041763 ],\n",
       "         [ 0.5430426 ,  0.83970505],\n",
       "         ...,\n",
       "         [ 0.99351764, -0.11367732],\n",
       "         [ 0.8001608 , -0.5997855 ],\n",
       "         [-0.8337144 ,  0.55219585]]],\n",
       "\n",
       "\n",
       "       [[[-0.6251835 ,  0.78047776],\n",
       "         [ 0.6738433 ,  0.7388742 ],\n",
       "         [-0.9261264 , -0.37721333],\n",
       "         ...,\n",
       "         [ 0.77851856, -0.6276216 ],\n",
       "         [ 0.6481785 ,  0.76148844],\n",
       "         [ 0.8758924 , -0.4825066 ]],\n",
       "\n",
       "        [[-0.9522389 , -0.3053537 ],\n",
       "         [-0.94333917, -0.33183014],\n",
       "         [-0.8840897 ,  0.46731716],\n",
       "         ...,\n",
       "         [ 0.8191087 ,  0.5736384 ],\n",
       "         [ 0.6949664 ,  0.7190423 ],\n",
       "         [-0.2683343 ,  0.96332586]],\n",
       "\n",
       "        [[-0.83271   ,  0.5537092 ],\n",
       "         [-0.72585243,  0.6878505 ],\n",
       "         [ 0.77075607,  0.63713026],\n",
       "         ...,\n",
       "         [-0.9409917 ,  0.33842987],\n",
       "         [-0.9817458 ,  0.19019753],\n",
       "         [ 0.9783392 ,  0.20700799]],\n",
       "\n",
       "        ...,\n",
       "\n",
       "        [[-0.04849374,  0.99882346],\n",
       "         [ 0.9695309 ,  0.24496901],\n",
       "         [-0.37880003, -0.9254785 ],\n",
       "         ...,\n",
       "         [ 0.6892487 , -0.7245249 ],\n",
       "         [ 0.81647533,  0.57738036],\n",
       "         [-0.9972612 , -0.07396041]],\n",
       "\n",
       "        [[-0.7055271 ,  0.7086829 ],\n",
       "         [-0.8819862 ,  0.4712752 ],\n",
       "         [-0.9642273 , -0.2650766 ],\n",
       "         ...,\n",
       "         [ 0.82875234, -0.55961555],\n",
       "         [-0.06711138,  0.9977455 ],\n",
       "         [-0.1260777 ,  0.99202037]],\n",
       "\n",
       "        [[ 0.99550056, -0.09475542],\n",
       "         [-0.9923597 ,  0.1233783 ],\n",
       "         [-0.71074235, -0.70345235],\n",
       "         ...,\n",
       "         [ 0.8930231 , -0.4500106 ],\n",
       "         [ 0.1744522 ,  0.9846656 ],\n",
       "         [ 0.8210883 , -0.570801  ]]],\n",
       "\n",
       "\n",
       "       ...,\n",
       "\n",
       "\n",
       "       [[[-0.7103158 ,  0.703883  ],\n",
       "         [ 0.37816712, -0.9257372 ],\n",
       "         [ 0.9501393 ,  0.31182557],\n",
       "         ...,\n",
       "         [ 0.66041577,  0.75090003],\n",
       "         [ 0.00253691,  0.9999967 ],\n",
       "         [-0.3108576 , -0.95045644]],\n",
       "\n",
       "        [[-0.9948109 ,  0.1017414 ],\n",
       "         [ 0.8379823 ,  0.54569715],\n",
       "         [-0.42345405,  0.9059175 ],\n",
       "         ...,\n",
       "         [ 0.7350055 ,  0.678061  ],\n",
       "         [-0.6902005 , -0.7236182 ],\n",
       "         [ 0.3050526 , -0.9523355 ]],\n",
       "\n",
       "        [[-0.9992436 , -0.03888722],\n",
       "         [-0.9524134 , -0.3048092 ],\n",
       "         [ 0.17434528,  0.9846846 ],\n",
       "         ...,\n",
       "         [-0.48622662, -0.8738327 ],\n",
       "         [ 0.5717803 ,  0.82040673],\n",
       "         [-0.7644213 ,  0.64471704]],\n",
       "\n",
       "        ...,\n",
       "\n",
       "        [[-0.0501435 , -0.998742  ],\n",
       "         [-0.915064  ,  0.40330857],\n",
       "         [-0.7114537 , -0.7027328 ],\n",
       "         ...,\n",
       "         [ 0.9969975 , -0.07743412],\n",
       "         [ 0.58687896, -0.80967456],\n",
       "         [ 0.9493494 ,  0.31422222]],\n",
       "\n",
       "        [[ 0.49507466,  0.86885047],\n",
       "         [ 0.8797779 , -0.47538495],\n",
       "         [-0.9211454 ,  0.38921857],\n",
       "         ...,\n",
       "         [ 0.8058742 , -0.5920867 ],\n",
       "         [ 0.53095114, -0.84740245],\n",
       "         [ 0.9835844 , -0.18044874]],\n",
       "\n",
       "        [[-0.19601943, -0.98059994],\n",
       "         [ 0.5338162 , -0.8456005 ],\n",
       "         [-0.9987141 , -0.05069751],\n",
       "         ...,\n",
       "         [ 0.73851573, -0.6742363 ],\n",
       "         [-0.99623436,  0.08670113],\n",
       "         [ 0.62578243,  0.7799976 ]]],\n",
       "\n",
       "\n",
       "       [[[ 0.99914247, -0.04140271],\n",
       "         [-0.45437112,  0.8908124 ],\n",
       "         [-0.42213327, -0.9065337 ],\n",
       "         ...,\n",
       "         [ 0.62287563,  0.78232086],\n",
       "         [-0.34597075, -0.9382453 ],\n",
       "         [ 0.79978997, -0.6002798 ]],\n",
       "\n",
       "        [[ 0.8313185 , -0.5557963 ],\n",
       "         [ 0.7764067 , -0.6302323 ],\n",
       "         [-0.4633804 ,  0.8861594 ],\n",
       "         ...,\n",
       "         [-0.06348424, -0.99798286],\n",
       "         [-0.24774338,  0.9688258 ],\n",
       "         [ 0.7214034 ,  0.6925151 ]],\n",
       "\n",
       "        [[-0.7895615 , -0.6136714 ],\n",
       "         [ 0.7669125 , -0.64175165],\n",
       "         [-0.60046875,  0.7996482 ],\n",
       "         ...,\n",
       "         [-0.22642732,  0.974028  ],\n",
       "         [-0.874757  , -0.48456183],\n",
       "         [-0.75133216, -0.6599242 ]],\n",
       "\n",
       "        ...,\n",
       "\n",
       "        [[ 0.99788564,  0.06499327],\n",
       "         [-0.9613605 ,  0.27529252],\n",
       "         [ 0.8874723 , -0.46086106],\n",
       "         ...,\n",
       "         [ 0.24368125,  0.96985537],\n",
       "         [ 0.9091159 , -0.4165433 ],\n",
       "         [ 0.39610097,  0.9182069 ]],\n",
       "\n",
       "        [[ 0.15766756, -0.98749226],\n",
       "         [-0.40235043,  0.9154857 ],\n",
       "         [-0.7679899 , -0.64046216],\n",
       "         ...,\n",
       "         [ 0.85627615,  0.5165182 ],\n",
       "         [ 0.43845195,  0.89875454],\n",
       "         [-0.995217  ,  0.09768905]],\n",
       "\n",
       "        [[ 0.31384754, -0.94947344],\n",
       "         [-0.335953  ,  0.94187874],\n",
       "         [ 0.32295498,  0.94641435],\n",
       "         ...,\n",
       "         [ 0.56830823,  0.82281566],\n",
       "         [ 0.5089668 ,  0.860786  ],\n",
       "         [ 0.43783924,  0.89905316]]],\n",
       "\n",
       "\n",
       "       [[[ 0.00368896, -0.99999326],\n",
       "         [-0.15407981, -0.98805845],\n",
       "         [-0.9993973 ,  0.03471346],\n",
       "         ...,\n",
       "         [-0.7566905 ,  0.6537732 ],\n",
       "         [-0.58953756,  0.80774087],\n",
       "         [-0.68740994,  0.72626966]],\n",
       "\n",
       "        [[-0.46347883,  0.8861079 ],\n",
       "         [ 0.00387764,  0.99999243],\n",
       "         [ 0.58661395, -0.80986667],\n",
       "         ...,\n",
       "         [ 0.9909095 ,  0.13453025],\n",
       "         [-0.44597933,  0.89504325],\n",
       "         [ 0.6099662 ,  0.7924275 ]],\n",
       "\n",
       "        [[ 0.88540715,  0.46481633],\n",
       "         [-0.9495712 ,  0.31355175],\n",
       "         [-0.9835763 , -0.1804929 ],\n",
       "         ...,\n",
       "         [ 0.11676199,  0.99315995],\n",
       "         [ 0.78835183, -0.6152245 ],\n",
       "         [ 0.6080205 , -0.7939213 ]],\n",
       "\n",
       "        ...,\n",
       "\n",
       "        [[ 0.8907368 , -0.45451924],\n",
       "         [-0.61389214, -0.7893899 ],\n",
       "         [-0.9944057 , -0.10562801],\n",
       "         ...,\n",
       "         [ 0.37469187,  0.92714936],\n",
       "         [ 0.94401246, -0.32990944],\n",
       "         [ 0.7751419 ,  0.6317871 ]],\n",
       "\n",
       "        [[-0.14082797, -0.99003404],\n",
       "         [ 0.27515417,  0.96140003],\n",
       "         [-0.40955904,  0.9122836 ],\n",
       "         ...,\n",
       "         [ 0.8188169 ,  0.5740549 ],\n",
       "         [-0.9845477 , -0.17511682],\n",
       "         [-0.71323305,  0.7009269 ]],\n",
       "\n",
       "        [[-0.9327261 , -0.36058545],\n",
       "         [-0.2761146 , -0.9611246 ],\n",
       "         [ 0.7902948 , -0.6127267 ],\n",
       "         ...,\n",
       "         [ 0.999295  ,  0.03754444],\n",
       "         [ 0.95306456, -0.30276698],\n",
       "         [-0.9659615 ,  0.2586859 ]]]], dtype=float32)>"
      ]
     },
     "execution_count": 111,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "r = tf.random.normal(shape = (16, 8, 8, 2))\n",
    "tf.linalg.l2_normalize(r,axis=-1)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".patch_gen_venv",
   "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.10.9"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}