batch_layer_norm.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Implementation of BatchNorm1D and LayerNorm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Intel MKL WARNING: Support of Intel(R) Streaming SIMD Extensions 4.2 (Intel(R) SSE4.2) enabled only processors has been deprecated. Intel oneAPI Math Kernel Library 2025.0 will require Intel(R) Advanced Vector Extensions (Intel(R) AVX) instructions.\n",
      "Intel MKL WARNING: Support of Intel(R) Streaming SIMD Extensions 4.2 (Intel(R) SSE4.2) enabled only processors has been deprecated. Intel oneAPI Math Kernel Library 2025.0 will require Intel(R) Advanced Vector Extensions (Intel(R) AVX) instructions.\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import torch.nn as nn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[7., 6.],\n",
       "         [8., 5.],\n",
       "         [1., 3.]],\n",
       "\n",
       "        [[8., 6.],\n",
       "         [5., 3.],\n",
       "         [5., 5.]]])"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.manual_seed(42)\n",
    "B,T,C = 2,3,2 # batch, time, channels\n",
    "x = torch.randint(1,10,(B,T,C)).float()\n",
    "x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### *BatchNorm1D*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[[ 0.54944179,  1.06904153],\n",
       "        [ 0.96152313,  0.26726038],\n",
       "        [-1.92304626, -1.33630191]],\n",
       "\n",
       "       [[ 0.96152313,  1.06904153],\n",
       "        [-0.27472089, -1.33630191],\n",
       "        [-0.27472089,  0.26726038]]])"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import numpy as np\n",
    "\n",
    "x = [[[7., 6.],\n",
    "      [8., 5.],\n",
    "      [1., 3.]],\n",
    "\n",
    "     [[8., 6.],\n",
    "      [5., 3.],\n",
    "      [5., 5.]]]\n",
    "\n",
    "feature_1 = np.array([7,8,1,8,5,5])\n",
    "feature_2 = np.array([6,5,3,6,3,5])\n",
    "\n",
    "feature_1_mean = np.mean(feature_1)\n",
    "feature_2_mean = np.mean(feature_2)\n",
    "\n",
    "feature_1_var = np.var(feature_1)\n",
    "feature_2_var = np.var(feature_2)\n",
    "\n",
    "means = np.array([[feature_1_mean, feature_2_mean]])\n",
    "vars = np.array([[feature_1_var, feature_2_var]])\n",
    "\n",
    "y_nonscaled = (x - means) / np.sqrt(vars + 1e-5)\n",
    "y_nonscaled\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "feature_1 = [7,8,1,8,5,5]\n",
    "feature_2 = [6,5,3,6,3,5]\n",
    "\n",
    "feature_1_mean = 5.67\n",
    "feature_2_mean = 4.67\n",
    "\n",
    "feature_1_std = 5.67\n",
    "feature_1_std = 5.67\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[ 0.5494,  1.0690],\n",
       "         [ 0.9615,  0.2673],\n",
       "         [-1.9230, -1.3363]],\n",
       "\n",
       "        [[ 0.9615,  1.0690],\n",
       "         [-0.2747, -1.3363],\n",
       "         [-0.2747,  0.2673]]], grad_fn=<TransposeBackward0>)"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "batch_norm = nn.BatchNorm1d(num_features=x.size(2))\n",
    "x_transposed = x.transpose(1, 2)\n",
    "y_normalized = batch_norm(x_transposed)\n",
    "torch_y = y_normalized.transpose(1, 2)\n",
    "torch_y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[ 0.5494,  1.0690],\n",
       "         [ 0.9615,  0.2673],\n",
       "         [-1.9230, -1.3363]],\n",
       "\n",
       "        [[ 0.9615,  1.0690],\n",
       "         [-0.2747, -1.3363],\n",
       "         [-0.2747,  0.2673]]])"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "epsilon = 1e-5\n",
    "gamma = torch.ones(x.shape[2])\n",
    "beta = torch.zeros(x.shape[2])\n",
    "\n",
    "# Calculate mean and variance for batch normalization\n",
    "mean = x.mean(dim=(0, 1), keepdim=True)  # Mean across batches and time steps\n",
    "variance = x.var(dim=(0, 1), unbiased=False, keepdim=True)  # Variance across batches and time steps\n",
    "\n",
    "# Normalize\n",
    "y = (x - mean) / torch.sqrt(variance + epsilon) * gamma + beta\n",
    "y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.allclose(y, torch_y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.allclose(torch.tensor(y_nonscaled, dtype=torch.float), torch_y)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### *LayerNorm*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[ 1.0000, -1.0000],\n",
       "         [ 1.0000, -1.0000],\n",
       "         [-1.0000,  1.0000]],\n",
       "\n",
       "        [[ 1.0000, -1.0000],\n",
       "         [ 1.0000, -1.0000],\n",
       "         [ 0.0000,  0.0000]]], grad_fn=<NativeLayerNormBackward0>)"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "layer_norm = torch.nn.LayerNorm(normalized_shape=[2])\n",
    "# Calculate LayerNorm\n",
    "torch_normalized = layer_norm(x)\n",
    "torch_normalized"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[ 1.0000, -1.0000],\n",
       "         [ 1.0000, -1.0000],\n",
       "         [-1.0000,  1.0000]],\n",
       "\n",
       "        [[ 1.0000, -1.0000],\n",
       "         [ 1.0000, -1.0000],\n",
       "         [ 0.0000,  0.0000]]])"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def layer_norm(x):\n",
    "  mu = x.mean(dim=2, keepdim=True)\n",
    "  var = x.var(dim=2, unbiased=False, keepdim=True)\n",
    "  eps = 1e-5\n",
    "\n",
    "  y = (x-mu) / torch.sqrt(var + eps)\n",
    "  return y\n",
    "\n",
    "x_normalized = layer_norm(x)\n",
    "x_normalized"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[[ 0.99998   , -0.99998   ],\n",
       "        [ 0.99999778, -0.99999778],\n",
       "        [-0.999995  ,  0.999995  ]],\n",
       "\n",
       "       [[ 0.999995  , -0.999995  ],\n",
       "        [ 0.999995  , -0.999995  ],\n",
       "        [ 0.        ,  0.        ]]])"
      ]
     },
     "execution_count": 50,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import numpy as np\n",
    "\n",
    "x = [[[7., 6.],\n",
    "      [8., 5.],\n",
    "      [1., 3.]],\n",
    "\n",
    "     [[8., 6.],\n",
    "      [5., 3.],\n",
    "      [5., 5.]]]\n",
    "\n",
    "means = np.mean(x, axis=2, keepdims=True) \n",
    "vars = np.var(x, axis=2, keepdims=True) \n",
    "\n",
    "y_nonscaled = (x - means) / (np.sqrt(vars + 1e-5))\n",
    "y_nonscaled"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "all equal!\n"
     ]
    }
   ],
   "source": [
    "assert torch.allclose(x_normalized, torch_normalized)\n",
    "assert torch.allclose(torch.tensor(y_nonscaled, dtype=torch.float), torch_normalized)\n",
    "print('all equal!')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "bitnet",
   "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.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}