{
"cells": [
{
"cell_type": "markdown",
"id": "c3edcae5",
"metadata": {
"editable": true
},
"source": [
""
]
},
{
"cell_type": "markdown",
"id": "4102577e",
"metadata": {
"editable": true
},
"source": [
"# Clustering and Unsupervised Learning\n",
"\n",
"In general terms cluster analysis, or clustering, is the task of grouping a\n",
"data-set into different distinct categories based on some measure of equality of\n",
"the data. This measure is often referred to as a **metric** or **similarity\n",
"measure** in the literature (note: sometimes we deal with a **dissimilarity\n",
"measure** instead). Usually, these metrics are formulated as some kind of\n",
"distance function between points in a high-dimensional space.\n",
"\n",
"The simplest, and also the most\n",
"common is the **Euclidean distance**.\n",
"\n",
"The simplest of all clustering algorithms is the **k-means algorithm**\n",
", sometimes also referred to as *Lloyds algorithm*. It is the simplest and also\n",
"the most common. From its simplicity it obtains both strengths and weaknesses.\n",
"These will be discussed in more detail later. The $k$-means algorithm is a\n",
"**centroid based** clustering algorithm.\n",
"\n",
"Assume, we are given $n$ data points and we wish to split the data into $K < n$\n",
"different categories, or clusters. We label each cluster by an integer"
]
},
{
"cell_type": "markdown",
"id": "0deb3255",
"metadata": {
"editable": true
},
"source": [
"$$\n",
"k\\in\\{1, \\cdots, K \\}.\n",
"$$"
]
},
{
"cell_type": "markdown",
"id": "cfd8fe00",
"metadata": {
"editable": true
},
"source": [
"In the basic k-means algorithm each point is assigned to only\n",
"one cluster $k$, and these assignments are *non-injective* i.e. many-to-one. We\n",
"can think of these mappings as an encoder $k = C(i)$, which assigns the $i$-th\n",
"data-point $\\bf x_i$ to the $k$-th cluster.\n",
"\n",
"$k$-means algorithm in words:\n",
"1. We start with guesses / random initializations of our $k$ cluster centers/centroids\n",
"\n",
"2. For each centroid the points that are most similar are identified\n",
"\n",
"3. Then we move / replace each centroid with a coordinate average of all the points that were assigned to that centroid.\n",
"\n",
"4. Iterate 2-3 until the centroids no longer move (to some tolerance)\n",
"\n",
"We assume we have $n$ data-points"
]
},
{
"cell_type": "markdown",
"id": "a29b7459",
"metadata": {
"editable": true
},
"source": [
"\n",
"\n",
"\n",
"$$\n",
"\\begin{equation}\\label{eq:kmeanspoints} \\tag{1}\n",
" \\boldsymbol{x_i} = \\{x_{i, 1}, \\cdots, x_{i, p}\\}\\in\\mathbb{R}^p.\n",
"\\end{equation}\n",
"$$"
]
},
{
"cell_type": "markdown",
"id": "98f9e37b",
"metadata": {
"editable": true
},
"source": [
"which we wish to group into $K < n$ clusters. For our dissimilarity measure we\n",
"use the *squared Euclidean distance*"
]
},
{
"cell_type": "markdown",
"id": "d3c32572",
"metadata": {
"editable": true
},
"source": [
"\n",
"\n",
"\n",
"$$\n",
"\\begin{equation}\\label{eq:squaredeuclidean} \\tag{2}\n",
" d(\\boldsymbol{x_i}, \\boldsymbol{x_i'}) = \\sum_{j=1}^p(x_{ij} - x_{i'j})^2\n",
" = ||\\boldsymbol{x_i} - \\boldsymbol{x_{i'}}||^2\n",
"\\end{equation}\n",
"$$"
]
},
{
"cell_type": "markdown",
"id": "29d24648",
"metadata": {
"editable": true
},
"source": [
"We define the so called *within-cluster point scatter* which gives us a\n",
"measure of how close each data point assigned to the same cluster tends to be to\n",
"the all the others."
]
},
{
"cell_type": "markdown",
"id": "fce5c797",
"metadata": {
"editable": true
},
"source": [
"\n",
"\n",
"\n",
"$$\n",
"\\begin{equation}\\label{eq:withincluster} \\tag{3}\n",
" W(C) = \\frac{1}{2}\\sum_{k=1}^K\\sum_{C(i)=k}\n",
" \\sum_{C(i')=k}d(\\boldsymbol{x_i}, \\boldsymbol{x_{i'}}) =\n",
" \\sum_{k=1}^KN_k\\sum_{C(i)=k}||\\boldsymbol{x_i} - \\boldsymbol{\\overline{x_k}}||^2\n",
"\\end{equation}\n",
"$$"
]
},
{
"cell_type": "markdown",
"id": "674a26b7",
"metadata": {
"editable": true
},
"source": [
"where $\\boldsymbol{\\overline{x_k}}$ is the mean vector associated with the $k$-th\n",
"cluster, and $N_k = \\sum_{i=1}^nI(C(i) = k)$, where the $I()$ notation is\n",
"similar to the Kronecker delta (*Commonly used in statistics, it just means that\n",
"when $i = k$ we have the encoder $C(i)$*). In other words, the within-cluster\n",
"scatter measures the compactness of each cluster with respect to the data points\n",
"assigned to each cluster. This is the quantity that the $k$-means algorithm aims\n",
"to minimize. We refer to this quantity $W(C)$ as the within cluster scatter\n",
"because of its relation to the *total scatter*.\n",
"\n",
"We have"
]
},
{
"cell_type": "markdown",
"id": "f200e7ff",
"metadata": {
"editable": true
},
"source": [
"\n",
"\n",
"\n",
"$$\n",
"\\begin{equation}\\label{eq:totalscatter} \\tag{4}\n",
" T = W(C) + B(C) = \\frac{1}{2}\\sum_{i=1}^n\n",
" \\sum_{i'=1}^nd(\\boldsymbol{x_i}, \\boldsymbol{x_{i'}})\n",
" = \\frac{1}{2}\\sum_{k=1}^K\\sum_{C(i)=k}\n",
" \\Big(\\sum_{C(i') = k}d(\\boldsymbol{x_i}, \\boldsymbol{x_{i'}})\n",
" + \\sum_{C(i')\\neq k}d(\\boldsymbol{x_i}, \\boldsymbol{x_{i'}})\\Big).\n",
"\\end{equation}\n",
"$$"
]
},
{
"cell_type": "markdown",
"id": "5471a94d",
"metadata": {
"editable": true
},
"source": [
"This is a quantity that is conserved throughout the $k$-means algorithm. It can\n",
"be thought of as the total amount of information in the data, and it is composed\n",
"of the aforementioned within-cluster scatter and the *between-cluster scatter*\n",
"$B(C)$. In methods such as principle component analysis the total scatter is not\n",
"conserved.\n",
"\n",
"Given a cluster mean $\\boldsymbol{m_k}$ we define the **total cluster variance**"
]
},
{
"cell_type": "markdown",
"id": "299a99ce",
"metadata": {
"editable": true
},
"source": [
"\n",
"\n",
"\n",
"$$\n",
"\\begin{equation}\\label{eq:totalclustervariance} \\tag{5}\n",
" \\min_{C, \\{\\boldsymbol{m_k}\\}_1^K}\\sum_{k=1}^KN_k\\sum||\\boldsymbol{x_i} - \\boldsymbol{m_k}||^2\n",
"\\end{equation}\n",
"$$"
]
},
{
"cell_type": "markdown",
"id": "bdfa54ee",
"metadata": {
"editable": true
},
"source": [
"Now we have all the pieces necessary to formally revisit the $k$-means algorithm.\n",
"\n",
"The $k$-means clustering algorithm goes as follows \n",
"\n",
"1. For a given cluster assignment $C$, and $k$ cluster means $\\left\\{m_1, \\cdots, m_k\\right\\}$. We minimize the total cluster variance with respect to the cluster means $\\{m_k\\}$ yielding the means of the currently assigned clusters.\n",
"\n",
"2. Given a current set of $k$ means $\\{m_k\\}$ the total cluster variance is minimized by assigning each observation to the closest (current) cluster mean. That is $$C(i) = \\underset{1\\leq k\\leq K}{\\mathrm{argmin}} ||\\boldsymbol{x_i} - \\boldsymbol{m_k}||^2$$\n",
"\n",
"3. Steps 1 and 2 are repeated until the assignments do not change."
]
},
{
"cell_type": "markdown",
"id": "f5def86c",
"metadata": {
"editable": true
},
"source": [
"## Codes and Approaches\n",
"\n",
"1. Before we start we specify a number $k$ which is the number of clusters we want to try to separate our data into.\n",
"\n",
"2. We initially choose $k$ random data points in our data as our initial centroids, *or means* (this is where the name comes from).\n",
"\n",
"3. Assign each data point to their closest centroid, based on the squared Euclidean distance.\n",
"\n",
"4. For each of the $k$ cluster we update the centroid by calculating new mean values for all the data points in the cluster.\n",
"\n",
"5. Iteratively minimize the within cluster scatter by performing steps (3, 4) until the new assignments stop changing (can be to some tolerance) or until a maximum number of iterations have passed.\n",
"\n",
"Let us now program the most basic version of the algorithm using nothing but\n",
"Python with numpy arrays. This code is kept intentionally simple to gradually\n",
"progress our understanding. There is no vectorization of any kind, and even most\n",
"helper functions are not utilized.\n",
"\n",
"We need first a dataset to do our cluster analysis on. In our case\n",
"this is a plain *vanilla* data set using random numbers using a\n",
"Gaussian distribution."
]
},
{
"cell_type": "code",
"execution_count": 1,
"id": "b0260188",
"metadata": {
"collapsed": false,
"editable": true
},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"/Users/mhjensen/miniforge3/envs/myenv/lib/python3.9/site-packages/jax/_src/lib/__init__.py:32: UserWarning: JAX on Mac ARM machines is experimental and minimally tested. Please see https://github.com/google/jax/issues/5501 in the event of problems.\n",
" warnings.warn(\"JAX on Mac ARM machines is experimental and minimally tested. \"\n"
]
},
{
"ename": "AttributeError",
"evalue": "module 'jaxlib.pocketfft' has no attribute 'pocketfft'",
"output_type": "error",
"traceback": [
"\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
"\u001b[0;31mAttributeError\u001b[0m Traceback (most recent call last)",
"Input \u001b[0;32mIn [1]\u001b[0m, in \u001b[0;36m| \u001b[0;34m()\u001b[0m\n\u001b[1;32m 3\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mtime\u001b[39;00m\n\u001b[1;32m 4\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mnumpy\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m \u001b[38;5;21;01mnp\u001b[39;00m\n\u001b[0;32m----> 5\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mtensorflow\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m \u001b[38;5;21;01mtf\u001b[39;00m\n\u001b[1;32m 6\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mmatplotlib\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m image\n\u001b[1;32m 7\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mmatplotlib\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mpyplot\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m \u001b[38;5;21;01mplt\u001b[39;00m\n",
"File \u001b[0;32m~/miniforge3/envs/myenv/lib/python3.9/site-packages/tensorflow/__init__.py:51\u001b[0m, in \u001b[0;36m\u001b[0;34m\u001b[0m\n\u001b[1;32m 49\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01m_api\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mv2\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m autograph\n\u001b[1;32m 50\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01m_api\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mv2\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m bitwise\n\u001b[0;32m---> 51\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01m_api\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mv2\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m compat\n\u001b[1;32m 52\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01m_api\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mv2\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m config\n\u001b[1;32m 53\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01m_api\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mv2\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m data\n",
"File \u001b[0;32m~/miniforge3/envs/myenv/lib/python3.9/site-packages/tensorflow/_api/v2/compat/__init__.py:37\u001b[0m, in \u001b[0;36m\u001b[0;34m\u001b[0m\n\u001b[1;32m 3\u001b[0m \u001b[38;5;124;03m\"\"\"Compatibility functions.\u001b[39;00m\n\u001b[1;32m 4\u001b[0m \n\u001b[1;32m 5\u001b[0m \u001b[38;5;124;03mThe `tf.compat` module contains two sets of compatibility functions.\u001b[39;00m\n\u001b[0;32m (...)\u001b[0m\n\u001b[1;32m 32\u001b[0m \n\u001b[1;32m 33\u001b[0m \u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m 35\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01msys\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m \u001b[38;5;21;01m_sys\u001b[39;00m\n\u001b[0;32m---> 37\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m v1\n\u001b[1;32m 38\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m v2\n\u001b[1;32m 39\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtensorflow\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mpython\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mcompat\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mcompat\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m forward_compatibility_horizon\n",
"File \u001b[0;32m~/miniforge3/envs/myenv/lib/python3.9/site-packages/tensorflow/_api/v2/compat/v1/__init__.py:30\u001b[0m, in \u001b[0;36m\u001b[0;34m\u001b[0m\n\u001b[1;32m 28\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m autograph\n\u001b[1;32m 29\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m bitwise\n\u001b[0;32m---> 30\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m compat\n\u001b[1;32m 31\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m config\n\u001b[1;32m 32\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m data\n",
"File \u001b[0;32m~/miniforge3/envs/myenv/lib/python3.9/site-packages/tensorflow/_api/v2/compat/v1/compat/__init__.py:37\u001b[0m, in \u001b[0;36m\u001b[0;34m\u001b[0m\n\u001b[1;32m 3\u001b[0m \u001b[38;5;124;03m\"\"\"Compatibility functions.\u001b[39;00m\n\u001b[1;32m 4\u001b[0m \n\u001b[1;32m 5\u001b[0m \u001b[38;5;124;03mThe `tf.compat` module contains two sets of compatibility functions.\u001b[39;00m\n\u001b[0;32m (...)\u001b[0m\n\u001b[1;32m 32\u001b[0m \n\u001b[1;32m 33\u001b[0m \u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m 35\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01msys\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m \u001b[38;5;21;01m_sys\u001b[39;00m\n\u001b[0;32m---> 37\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m v1\n\u001b[1;32m 38\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m v2\n\u001b[1;32m 39\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtensorflow\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mpython\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mcompat\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mcompat\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m forward_compatibility_horizon\n",
"File \u001b[0;32m~/miniforge3/envs/myenv/lib/python3.9/site-packages/tensorflow/_api/v2/compat/v1/compat/v1/__init__.py:47\u001b[0m, in \u001b[0;36m\u001b[0;34m\u001b[0m\n\u001b[1;32m 45\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtensorflow\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01m_api\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mv2\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mcompat\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mv1\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m layers\n\u001b[1;32m 46\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtensorflow\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01m_api\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mv2\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mcompat\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mv1\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m linalg\n\u001b[0;32m---> 47\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtensorflow\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01m_api\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mv2\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mcompat\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mv1\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m lite\n\u001b[1;32m 48\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtensorflow\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01m_api\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mv2\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mcompat\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mv1\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m logging\n\u001b[1;32m 49\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtensorflow\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01m_api\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mv2\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mcompat\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mv1\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m lookup\n",
"File \u001b[0;32m~/miniforge3/envs/myenv/lib/python3.9/site-packages/tensorflow/_api/v2/compat/v1/lite/__init__.py:9\u001b[0m, in \u001b[0;36m\u001b[0;34m\u001b[0m\n\u001b[1;32m 6\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01msys\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m \u001b[38;5;21;01m_sys\u001b[39;00m\n\u001b[1;32m 8\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m constants\n\u001b[0;32m----> 9\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m experimental\n\u001b[1;32m 10\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtensorflow\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mlite\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mpython\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mlite\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m Interpreter\n\u001b[1;32m 11\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtensorflow\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mlite\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mpython\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mlite\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m OpHint\n",
"File \u001b[0;32m~/miniforge3/envs/myenv/lib/python3.9/site-packages/tensorflow/_api/v2/compat/v1/lite/experimental/__init__.py:8\u001b[0m, in \u001b[0;36m\u001b[0;34m\u001b[0m\n\u001b[1;32m 3\u001b[0m \u001b[38;5;124;03m\"\"\"Public API for tf.lite.experimental namespace.\u001b[39;00m\n\u001b[1;32m 4\u001b[0m \u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m 6\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01msys\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m \u001b[38;5;21;01m_sys\u001b[39;00m\n\u001b[0;32m----> 8\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m authoring\n\u001b[1;32m 9\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtensorflow\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mlite\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mpython\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01manalyzer\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m ModelAnalyzer \u001b[38;5;28;01mas\u001b[39;00m Analyzer\n\u001b[1;32m 10\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtensorflow\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mlite\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mpython\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mlite\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m OpResolverType\n",
"File \u001b[0;32m~/miniforge3/envs/myenv/lib/python3.9/site-packages/tensorflow/_api/v2/compat/v1/lite/experimental/authoring/__init__.py:8\u001b[0m, in \u001b[0;36m\u001b[0;34m\u001b[0m\n\u001b[1;32m 3\u001b[0m \u001b[38;5;124;03m\"\"\"Public API for tf.lite.experimental.authoring namespace.\u001b[39;00m\n\u001b[1;32m 4\u001b[0m \u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m 6\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01msys\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m \u001b[38;5;21;01m_sys\u001b[39;00m\n\u001b[0;32m----> 8\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtensorflow\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mlite\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mpython\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mauthoring\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mauthoring\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m compatible\n",
"File \u001b[0;32m~/miniforge3/envs/myenv/lib/python3.9/site-packages/tensorflow/lite/python/authoring/authoring.py:43\u001b[0m, in \u001b[0;36m\u001b[0;34m\u001b[0m\n\u001b[1;32m 39\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mfunctools\u001b[39;00m\n\u001b[1;32m 42\u001b[0m \u001b[38;5;66;03m# pylint: disable=g-import-not-at-top\u001b[39;00m\n\u001b[0;32m---> 43\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtensorflow\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mlite\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mpython\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m convert\n\u001b[1;32m 44\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtensorflow\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mlite\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mpython\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m lite\n\u001b[1;32m 45\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtensorflow\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mlite\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mpython\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mmetrics\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m converter_error_data_pb2\n",
"File \u001b[0;32m~/miniforge3/envs/myenv/lib/python3.9/site-packages/tensorflow/lite/python/convert.py:29\u001b[0m, in \u001b[0;36m\u001b[0;34m\u001b[0m\n\u001b[1;32m 26\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01msix\u001b[39;00m\n\u001b[1;32m 28\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtensorflow\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mlite\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mpython\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m lite_constants\n\u001b[0;32m---> 29\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtensorflow\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mlite\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mpython\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m util\n\u001b[1;32m 30\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtensorflow\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mlite\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mpython\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m wrap_toco\n\u001b[1;32m 31\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtensorflow\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mlite\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mpython\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mconvert_phase\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m Component\n",
"File \u001b[0;32m~/miniforge3/envs/myenv/lib/python3.9/site-packages/tensorflow/lite/python/util.py:51\u001b[0m, in \u001b[0;36m\u001b[0;34m\u001b[0m\n\u001b[1;32m 47\u001b[0m \u001b[38;5;66;03m# Jax functions used by TFLite\u001b[39;00m\n\u001b[1;32m 48\u001b[0m \u001b[38;5;66;03m# pylint: disable=g-import-not-at-top\u001b[39;00m\n\u001b[1;32m 49\u001b[0m \u001b[38;5;66;03m# pylint: disable=unused-import\u001b[39;00m\n\u001b[1;32m 50\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m---> 51\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mjax\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m xla_computation \u001b[38;5;28;01mas\u001b[39;00m _xla_computation\n\u001b[1;32m 52\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mImportError\u001b[39;00m:\n\u001b[1;32m 53\u001b[0m _xla_computation \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n",
"File \u001b[0;32m~/miniforge3/envs/myenv/lib/python3.9/site-packages/jax/__init__.py:116\u001b[0m, in \u001b[0;36m\u001b[0;34m\u001b[0m\n\u001b[1;32m 40\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01m_src\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mconfig\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m (\n\u001b[1;32m 41\u001b[0m config \u001b[38;5;28;01mas\u001b[39;00m config,\n\u001b[1;32m 42\u001b[0m enable_checks \u001b[38;5;28;01mas\u001b[39;00m enable_checks,\n\u001b[0;32m (...)\u001b[0m\n\u001b[1;32m 51\u001b[0m numpy_rank_promotion \u001b[38;5;28;01mas\u001b[39;00m numpy_rank_promotion,\n\u001b[1;32m 52\u001b[0m )\n\u001b[1;32m 53\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01m_src\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mapi\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m (\n\u001b[1;32m 54\u001b[0m ad, \u001b[38;5;66;03m# TODO(phawkins): update users to avoid this.\u001b[39;00m\n\u001b[1;32m 55\u001b[0m checkpoint \u001b[38;5;28;01mas\u001b[39;00m checkpoint,\n\u001b[0;32m (...)\u001b[0m\n\u001b[1;32m 114\u001b[0m xla_computation \u001b[38;5;28;01mas\u001b[39;00m xla_computation,\n\u001b[1;32m 115\u001b[0m )\n\u001b[0;32m--> 116\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mexperimental\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mmaps\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m soft_pmap \u001b[38;5;28;01mas\u001b[39;00m soft_pmap\n\u001b[1;32m 117\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mversion\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m __version__ \u001b[38;5;28;01mas\u001b[39;00m __version__\n\u001b[1;32m 119\u001b[0m \u001b[38;5;66;03m# These submodules are separate because they are in an import cycle with\u001b[39;00m\n\u001b[1;32m 120\u001b[0m \u001b[38;5;66;03m# jax and rely on the names imported above.\u001b[39;00m\n",
"File \u001b[0;32m~/miniforge3/envs/myenv/lib/python3.9/site-packages/jax/experimental/maps.py:26\u001b[0m, in \u001b[0;36m\u001b[0;34m\u001b[0m\n\u001b[1;32m 23\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mfunctools\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m wraps, partial, partialmethod\n\u001b[1;32m 24\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01menum\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m Enum\n\u001b[0;32m---> 26\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m numpy \u001b[38;5;28;01mas\u001b[39;00m jnp\n\u001b[1;32m 27\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m core\n\u001b[1;32m 28\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m linear_util \u001b[38;5;28;01mas\u001b[39;00m lu\n",
"File \u001b[0;32m~/miniforge3/envs/myenv/lib/python3.9/site-packages/jax/numpy/__init__.py:19\u001b[0m, in \u001b[0;36m\u001b[0;34m\u001b[0m\n\u001b[1;32m 1\u001b[0m \u001b[38;5;66;03m# Copyright 2018 Google LLC\u001b[39;00m\n\u001b[1;32m 2\u001b[0m \u001b[38;5;66;03m#\u001b[39;00m\n\u001b[1;32m 3\u001b[0m \u001b[38;5;66;03m# Licensed under the Apache License, Version 2.0 (the \"License\");\u001b[39;00m\n\u001b[0;32m (...)\u001b[0m\n\u001b[1;32m 17\u001b[0m \n\u001b[1;32m 18\u001b[0m \u001b[38;5;66;03m# flake8: noqa: F401\u001b[39;00m\n\u001b[0;32m---> 19\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m fft \u001b[38;5;28;01mas\u001b[39;00m fft\n\u001b[1;32m 20\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m linalg \u001b[38;5;28;01mas\u001b[39;00m linalg\n\u001b[1;32m 22\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mjax\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01minterpreters\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mxla\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m DeviceArray \u001b[38;5;28;01mas\u001b[39;00m DeviceArray\n",
"File \u001b[0;32m~/miniforge3/envs/myenv/lib/python3.9/site-packages/jax/numpy/fft.py:17\u001b[0m, in \u001b[0;36m\u001b[0;34m\u001b[0m\n\u001b[1;32m 1\u001b[0m \u001b[38;5;66;03m# Copyright 2020 Google LLC\u001b[39;00m\n\u001b[1;32m 2\u001b[0m \u001b[38;5;66;03m#\u001b[39;00m\n\u001b[1;32m 3\u001b[0m \u001b[38;5;66;03m# Licensed under the Apache License, Version 2.0 (the \"License\");\u001b[39;00m\n\u001b[0;32m (...)\u001b[0m\n\u001b[1;32m 14\u001b[0m \n\u001b[1;32m 15\u001b[0m \u001b[38;5;66;03m# flake8: noqa: F401\u001b[39;00m\n\u001b[0;32m---> 17\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mjax\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01m_src\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mnumpy\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mfft\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m (\n\u001b[1;32m 18\u001b[0m ifft \u001b[38;5;28;01mas\u001b[39;00m ifft,\n\u001b[1;32m 19\u001b[0m ifft2 \u001b[38;5;28;01mas\u001b[39;00m ifft2,\n\u001b[1;32m 20\u001b[0m ifftn \u001b[38;5;28;01mas\u001b[39;00m ifftn,\n\u001b[1;32m 21\u001b[0m ifftshift \u001b[38;5;28;01mas\u001b[39;00m ifftshift,\n\u001b[1;32m 22\u001b[0m ihfft \u001b[38;5;28;01mas\u001b[39;00m ihfft,\n\u001b[1;32m 23\u001b[0m irfft \u001b[38;5;28;01mas\u001b[39;00m irfft,\n\u001b[1;32m 24\u001b[0m irfft2 \u001b[38;5;28;01mas\u001b[39;00m irfft2,\n\u001b[1;32m 25\u001b[0m irfftn \u001b[38;5;28;01mas\u001b[39;00m irfftn,\n\u001b[1;32m 26\u001b[0m fft \u001b[38;5;28;01mas\u001b[39;00m fft,\n\u001b[1;32m 27\u001b[0m fft2 \u001b[38;5;28;01mas\u001b[39;00m fft2,\n\u001b[1;32m 28\u001b[0m fftfreq \u001b[38;5;28;01mas\u001b[39;00m fftfreq,\n\u001b[1;32m 29\u001b[0m fftn \u001b[38;5;28;01mas\u001b[39;00m fftn,\n\u001b[1;32m 30\u001b[0m fftshift \u001b[38;5;28;01mas\u001b[39;00m fftshift,\n\u001b[1;32m 31\u001b[0m hfft \u001b[38;5;28;01mas\u001b[39;00m hfft,\n\u001b[1;32m 32\u001b[0m rfft \u001b[38;5;28;01mas\u001b[39;00m rfft,\n\u001b[1;32m 33\u001b[0m rfft2 \u001b[38;5;28;01mas\u001b[39;00m rfft2,\n\u001b[1;32m 34\u001b[0m rfftfreq \u001b[38;5;28;01mas\u001b[39;00m rfftfreq,\n\u001b[1;32m 35\u001b[0m rfftn \u001b[38;5;28;01mas\u001b[39;00m rfftn,\n\u001b[1;32m 36\u001b[0m )\n\u001b[1;32m 38\u001b[0m \u001b[38;5;66;03m# Module initialization is encapsulated in a function to avoid accidental\u001b[39;00m\n\u001b[1;32m 39\u001b[0m \u001b[38;5;66;03m# namespace pollution.\u001b[39;00m\n\u001b[1;32m 40\u001b[0m _NOT_IMPLEMENTED \u001b[38;5;241m=\u001b[39m []\n",
"File \u001b[0;32m~/miniforge3/envs/myenv/lib/python3.9/site-packages/jax/_src/numpy/fft.py:19\u001b[0m, in \u001b[0;36m\u001b[0;34m\u001b[0m\n\u001b[1;32m 16\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01moperator\u001b[39;00m\n\u001b[1;32m 17\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mnumpy\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m \u001b[38;5;21;01mnp\u001b[39;00m\n\u001b[0;32m---> 19\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mjax\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m lax\n\u001b[1;32m 20\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mjax\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01m_src\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mlib\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m xla_client\n\u001b[1;32m 21\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mjax\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01m_src\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mutil\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m safe_zip\n",
"File \u001b[0;32m~/miniforge3/envs/myenv/lib/python3.9/site-packages/jax/lax/__init__.py:332\u001b[0m, in \u001b[0;36m\u001b[0;34m\u001b[0m\n\u001b[1;32m 299\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mjax\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01m_src\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mlax\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mlax\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m (_reduce_sum, _reduce_max, _reduce_min, _reduce_or,\n\u001b[1;32m 300\u001b[0m _reduce_and, _reduce_window_sum, _reduce_window_max,\n\u001b[1;32m 301\u001b[0m _reduce_window_min, _reduce_window_prod,\n\u001b[0;32m (...)\u001b[0m\n\u001b[1;32m 306\u001b[0m _upcast_fp16_for_computation, _broadcasting_shape_rule,\n\u001b[1;32m 307\u001b[0m _eye, _tri, _delta, _ones, _zeros, _dilate_shape)\n\u001b[1;32m 308\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mjax\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01m_src\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mlax\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mcontrol_flow\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m (\n\u001b[1;32m 309\u001b[0m associative_scan \u001b[38;5;28;01mas\u001b[39;00m associative_scan,\n\u001b[1;32m 310\u001b[0m cond \u001b[38;5;28;01mas\u001b[39;00m cond,\n\u001b[0;32m (...)\u001b[0m\n\u001b[1;32m 330\u001b[0m while_p \u001b[38;5;28;01mas\u001b[39;00m while_p,\n\u001b[1;32m 331\u001b[0m )\n\u001b[0;32m--> 332\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mjax\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01m_src\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mlax\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mfft\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m (\n\u001b[1;32m 333\u001b[0m fft \u001b[38;5;28;01mas\u001b[39;00m fft,\n\u001b[1;32m 334\u001b[0m fft_p \u001b[38;5;28;01mas\u001b[39;00m fft_p,\n\u001b[1;32m 335\u001b[0m )\n\u001b[1;32m 336\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mjax\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01m_src\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mlax\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mparallel\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m (\n\u001b[1;32m 337\u001b[0m all_gather \u001b[38;5;28;01mas\u001b[39;00m all_gather,\n\u001b[1;32m 338\u001b[0m all_to_all \u001b[38;5;28;01mas\u001b[39;00m all_to_all,\n\u001b[0;32m (...)\u001b[0m\n\u001b[1;32m 355\u001b[0m xeinsum \u001b[38;5;28;01mas\u001b[39;00m xeinsum,\n\u001b[1;32m 356\u001b[0m )\n\u001b[1;32m 357\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mjax\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01m_src\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mlax\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mother\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m (\n\u001b[1;32m 358\u001b[0m conv_general_dilated_patches \u001b[38;5;28;01mas\u001b[39;00m conv_general_dilated_patches\n\u001b[1;32m 359\u001b[0m )\n",
"File \u001b[0;32m~/miniforge3/envs/myenv/lib/python3.9/site-packages/jax/_src/lax/fft.py:145\u001b[0m, in \u001b[0;36m\u001b[0;34m\u001b[0m\n\u001b[1;32m 143\u001b[0m batching\u001b[38;5;241m.\u001b[39mprimitive_batchers[fft_p] \u001b[38;5;241m=\u001b[39m fft_batching_rule\n\u001b[1;32m 144\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m pocketfft:\n\u001b[0;32m--> 145\u001b[0m xla\u001b[38;5;241m.\u001b[39mbackend_specific_translations[\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mcpu\u001b[39m\u001b[38;5;124m'\u001b[39m][fft_p] \u001b[38;5;241m=\u001b[39m \u001b[43mpocketfft\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mpocketfft\u001b[49m\n",
"\u001b[0;31mAttributeError\u001b[0m: module 'jaxlib.pocketfft' has no attribute 'pocketfft'"
]
}
],
"source": [
"%matplotlib inline\n",
"\n",
"import time\n",
"import numpy as np\n",
"import tensorflow as tf\n",
"from matplotlib import image\n",
"import matplotlib.pyplot as plt\n",
"from sklearn.cluster import KMeans\n",
"from IPython.display import display\n",
"\n",
"np.random.seed(2021)"
]
},
{
"cell_type": "markdown",
"id": "fe680e35",
"metadata": {
"editable": true
},
"source": [
"Next we define functions, for ease of use later, to generate Gaussians and to\n",
"set up our toy data set."
]
},
{
"cell_type": "code",
"execution_count": 2,
"id": "9db2bbce",
"metadata": {
"collapsed": false,
"editable": true
},
"outputs": [],
"source": [
"def gaussian_points(dim=2, n_points=1000, mean_vector=np.array([0, 0]),\n",
" sample_variance=1):\n",
" \"\"\"\n",
" Very simple custom function to generate gaussian distributed point clusters\n",
" with variable dimension, number of points, means in each direction\n",
" (must match dim) and sample variance.\n",
"\n",
" Inputs:\n",
" dim (int)\n",
" n_points (int)\n",
" mean_vector (np.array) (where index 0 is x, index 1 is y etc.)\n",
" sample_variance (float)\n",
"\n",
" Returns:\n",
" data (np.array): with dimensions (dim x n_points)\n",
" \"\"\"\n",
"\n",
" mean_matrix = np.zeros(dim) + mean_vector\n",
" covariance_matrix = np.eye(dim) * sample_variance\n",
" data = np.random.multivariate_normal(mean_matrix, covariance_matrix,\n",
" n_points)\n",
" return data\n",
"\n",
"\n",
"\n",
"def generate_simple_clustering_dataset(dim=2, n_points=1000, plotting=True,\n",
" return_data=True):\n",
" \"\"\"\n",
" Toy model to illustrate k-means clustering\n",
" \"\"\"\n",
"\n",
" data1 = gaussian_points(mean_vector=np.array([5, 5]))\n",
" data2 = gaussian_points()\n",
" data3 = gaussian_points(mean_vector=np.array([1, 4.5]))\n",
" data4 = gaussian_points(mean_vector=np.array([5, 1]))\n",
" data = np.concatenate((data1, data2, data3, data4), axis=0)\n",
"\n",
" if plotting:\n",
" fig, ax = plt.subplots()\n",
" ax.scatter(data[:, 0], data[:, 1], alpha=0.2)\n",
" ax.set_title('Toy Model Dataset')\n",
" plt.show()\n",
"\n",
"\n",
" if return_data:\n",
" return data\n",
"\n",
"\n",
"data = generate_simple_clustering_dataset()"
]
},
{
"cell_type": "markdown",
"id": "c0bb8c76",
"metadata": {
"editable": true
},
"source": [
"With the above dataset we start\n",
"implementing the $k$-means algorithm."
]
},
{
"cell_type": "code",
"execution_count": 3,
"id": "29a75065",
"metadata": {
"collapsed": false,
"editable": true
},
"outputs": [],
"source": [
"\n",
"n_samples, dimensions = data.shape\n",
"n_clusters = 4\n",
"\n",
"# we randomly initialize our centroids\n",
"np.random.seed(2021)\n",
"centroids = data[np.random.choice(n_samples, n_clusters, replace=False), :]\n",
"distances = np.zeros((n_samples, n_clusters))\n",
"\n",
"# first we need to calculate the distance to each centroid from our data\n",
"for k in range(n_clusters):\n",
" for n in range(n_samples):\n",
" dist = 0\n",
" for d in range(dimensions):\n",
" dist += np.abs(data[n, d] - centroids[k, d])**2\n",
" distances[n, k] = dist\n",
"\n",
"# we initialize an array to keep track of to which cluster each point belongs\n",
"# the way we set it up here the index tracks which point and the value which\n",
"# cluster the point belongs to\n",
"cluster_labels = np.zeros(n_samples, dtype='int')\n",
"\n",
"# next we loop through our samples and for every point assign it to the cluster\n",
"# to which it has the smallest distance to\n",
"for n in range(n_samples):\n",
" # tracking variables (all of this is basically just an argmin)\n",
" smallest = 1e10\n",
" smallest_row_index = 1e10\n",
" for k in range(n_clusters):\n",
" if distances[n, k] < smallest:\n",
" smallest = distances[n, k]\n",
" smallest_row_index = k\n",
"\n",
" cluster_labels[n] = smallest_row_index"
]
},
{
"cell_type": "code",
"execution_count": 4,
"id": "9fae7fc9",
"metadata": {
"collapsed": false,
"editable": true
},
"outputs": [],
"source": [
"fig = plt.figure()\n",
"ax = fig.add_subplot()\n",
"unique_cluster_labels = np.unique(cluster_labels)\n",
"for i in unique_cluster_labels:\n",
" ax.scatter(data[cluster_labels == i, 0],\n",
" data[cluster_labels == i, 1],\n",
" label = i,\n",
" alpha = 0.2)\n",
" ax.scatter(centroids[:, 0], centroids[:, 1], c='black')\n",
"\n",
"ax.set_title(\"First Grouping of Points to Centroids\")\n",
"\n",
"plt.show()"
]
},
{
"cell_type": "markdown",
"id": "90d2a873",
"metadata": {
"editable": true
},
"source": [
"So what do we have so far? We have 'picked' $k$ centroids at random from our\n",
"data points. There are other ways of more intelligently choosing their\n",
"initializations, however for our purposes randomly is fine. Then we have\n",
"initialized an array 'distances' which holds the information of the distance,\n",
"*or dissimilarity*, of every point to of our centroids. Finally, we have\n",
"initialized an array 'cluster_labels' which according to our distances array\n",
"holds the information of to which centroid every point is assigned. This was the\n",
"first pass of our algorithm. Essentially, all we need to do now is repeat the\n",
"distance and assignment steps above until we have reached a desired convergence\n",
"or a maximum amount of iterations."
]
},
{
"cell_type": "code",
"execution_count": 5,
"id": "378c29fc",
"metadata": {
"collapsed": false,
"editable": true
},
"outputs": [],
"source": [
"\n",
"max_iterations = 100\n",
"tolerance = 1e-8\n",
"\n",
"for iteration in range(max_iterations):\n",
" prev_centroids = centroids.copy()\n",
" for k in range(n_clusters):\n",
" # this array will be used to update our centroid positions\n",
" vector_mean = np.zeros(dimensions)\n",
" mean_divisor = 0\n",
" for n in range(n_samples):\n",
" if cluster_labels[n] == k:\n",
" vector_mean += data[n, :]\n",
" mean_divisor += 1\n",
"\n",
" # update according to the k means\n",
" centroids[k, :] = vector_mean / mean_divisor\n",
"\n",
" # we find the dissimilarity\n",
" for k in range(n_clusters):\n",
" for n in range(n_samples):\n",
" dist = 0\n",
" for d in range(dimensions):\n",
" dist += np.abs(data[n, d] - centroids[k, d])**2\n",
" distances[n, k] = dist\n",
"\n",
" # assign each point\n",
" for n in range(n_samples):\n",
" smallest = 1e10\n",
" smallest_row_index = 1e10\n",
" for k in range(n_clusters):\n",
" if distances[n, k] < smallest:\n",
" smallest = distances[n, k]\n",
" smallest_row_index = k\n",
"\n",
" cluster_labels[n] = smallest_row_index\n",
"\n",
" # convergence criteria\n",
" centroid_difference = np.sum(np.abs(centroids - prev_centroids))\n",
" if centroid_difference < tolerance:\n",
" print(f'Converged at iteration {iteration}')\n",
" break\n",
"\n",
" elif iteration == max_iterations:\n",
" print(f'Did not converge in {max_iterations} iterations')"
]
},
{
"cell_type": "markdown",
"id": "545a6742",
"metadata": {
"editable": true
},
"source": [
"We now have a simple , un-optimized $k$-means\n",
"clustering implementation. Lets plot the final result"
]
},
{
"cell_type": "code",
"execution_count": 6,
"id": "d9d3973b",
"metadata": {
"collapsed": false,
"editable": true
},
"outputs": [],
"source": [
"fig = plt.figure()\n",
"ax = fig.add_subplot()\n",
"unique_cluster_labels = np.unique(cluster_labels)\n",
"for i in unique_cluster_labels:\n",
" ax.scatter(data[cluster_labels == i, 0],\n",
" data[cluster_labels == i, 1],\n",
" label = i,\n",
" alpha = 0.2)\n",
" ax.scatter(centroids[:, 0], centroids[:, 1], c='black')\n",
"\n",
"ax.set_title(\"Final Result of K-means Clustering\")\n",
"\n",
"plt.show()"
]
},
{
"cell_type": "code",
"execution_count": 7,
"id": "ee6a145f",
"metadata": {
"collapsed": false,
"editable": true
},
"outputs": [],
"source": [
"def naive_kmeans(data, n_clusters=4, max_iterations=100, tolerance=1e-8):\n",
" start_time = time.time()\n",
"\n",
" n_samples, dimensions = data.shape\n",
" n_clusters = 4\n",
" #np.random.seed(2021)\n",
" centroids = data[np.random.choice(n_samples, n_clusters, replace=False), :]\n",
" distances = np.zeros((n_samples, n_clusters))\n",
"\n",
" for k in range(n_clusters):\n",
" for n in range(n_samples):\n",
" dist = 0\n",
" for d in range(dimensions):\n",
" dist += np.abs(data[n, d] - centroids[k, d])**2\n",
" distances[n, k] = dist\n",
"\n",
" cluster_labels = np.zeros(n_samples, dtype='int')\n",
"\n",
" for n in range(n_samples):\n",
" smallest = 1e10\n",
" smallest_row_index = 1e10\n",
" for k in range(n_clusters):\n",
" if distances[n, k] < smallest:\n",
" smallest = distances[n, k]\n",
" smallest_row_index = k\n",
"\n",
" cluster_labels[n] = smallest_row_index\n",
"\n",
" for iteration in range(max_iterations):\n",
" prev_centroids = centroids.copy()\n",
" for k in range(n_clusters):\n",
" vector_mean = np.zeros(dimensions)\n",
" mean_divisor = 0\n",
" for n in range(n_samples):\n",
" if cluster_labels[n] == k:\n",
" vector_mean += data[n, :]\n",
" mean_divisor += 1\n",
"\n",
" centroids[k, :] = vector_mean / mean_divisor\n",
"\n",
" for k in range(n_clusters):\n",
" for n in range(n_samples):\n",
" dist = 0\n",
" for d in range(dimensions):\n",
" dist += np.abs(data[n, d] - centroids[k, d])**2\n",
" distances[n, k] = dist\n",
"\n",
" for n in range(n_samples):\n",
" smallest = 1e10\n",
" smallest_row_index = 1e10\n",
" for k in range(n_clusters):\n",
" if distances[n, k] < smallest:\n",
" smallest = distances[n, k]\n",
" smallest_row_index = k\n",
"\n",
" cluster_labels[n] = smallest_row_index\n",
"\n",
" centroid_difference = np.sum(np.abs(centroids - prev_centroids))\n",
" if centroid_difference < tolerance:\n",
" print(f'Converged at iteration {iteration}')\n",
" print(f'Runtime: {time.time() - start_time} seconds')\n",
"\n",
" return cluster_labels, centroids\n",
"\n",
" print(f'Did not converge in {max_iterations} iterations')\n",
" print(f'Runtime: {time.time() - start_time} seconds')\n",
"\n",
" return cluster_labels, centroids"
]
}
],
"metadata": {
"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.9.10"
}
},
"nbformat": 4,
"nbformat_minor": 5
} |