{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Obtaining file:///Users/carlostrujillo/Documents/GitHub/pytensor\n",
            "  Installing build dependencies ... \u001b[?25ldone\n",
            "\u001b[?25h  Checking if build backend supports build_editable ... \u001b[?25ldone\n",
            "\u001b[?25h  Getting requirements to build editable ... \u001b[?25ldone\n",
            "\u001b[?25h  Preparing editable metadata (pyproject.toml) ... \u001b[?25ldone\n",
            "\u001b[?25hBuilding wheels for collected packages: pytensor\n",
            "  Building editable for pytensor (pyproject.toml) ... \u001b[?25ldone\n",
            "\u001b[?25h  Created wheel for pytensor: filename=pytensor-2.31.7+80.g06ccf91ba.dirty-0.editable-cp312-cp312-macosx_11_0_arm64.whl size=7323 sha256=c09587a5f3141d49000666d2817c5a01436f13ff5a19aa3deda20f647660afee\n",
            "  Stored in directory: /private/var/folders/f0/rbz8xs8s17n3k3f_ccp31bvh0000gn/T/pip-ephem-wheel-cache-i00nb67k/wheels/52/f6/4c/e6784e2203d5405c94db1d544248730e598e4397674416af05\n",
            "Successfully built pytensor\n",
            "Installing collected packages: pytensor\n",
            "  Attempting uninstall: pytensor\n",
            "    Found existing installation: pytensor 2.31.7+80.g06ccf91ba.dirty\n",
            "    Uninstalling pytensor-2.31.7+80.g06ccf91ba.dirty:\n",
            "      Successfully uninstalled pytensor-2.31.7+80.g06ccf91ba.dirty\n",
            "Successfully installed pytensor-2.31.7+80.g06ccf91ba.dirty\n",
            "Note: you may need to restart the kernel to use updated packages.\n"
          ]
        }
      ],
      "source": [
        "%pip install -e ../.. --no-deps"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {},
      "outputs": [],
      "source": [
        "import time\n",
        "import numpy as np\n",
        "import jax\n",
        "import jax.numpy as jnp\n",
        "\n",
        "import pytensor\n",
        "import pytensor.tensor as pt\n",
        "from pytensor.compile.function import function\n",
        "from pytensor.compile.mode import Mode\n",
        "from pytensor.graph import RewriteDatabaseQuery\n",
        "from pytensor.link.jax import JAXLinker\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {},
      "outputs": [],
      "source": [
        "# Configure JAX to use float32 for consistency with MLX\n",
        "jax.config.update(\"jax_enable_x64\", False)\n",
        "\n",
        "# Set up PyTensor JAX mode\n",
        "jax_optimizer = RewriteDatabaseQuery(include=[\"jax\"], exclude=[])\n",
        "pytensor_jax_mode = \"JAX\"\n",
        "\n",
        "# Try to set up MLX mode\n",
        "try:\n",
        "    from pytensor.link.mlx import MLXLinker\n",
        "    import mlx.core as mx\n",
        "    mlx_optimizer = RewriteDatabaseQuery(include=[\"mlx\"], exclude=[])\n",
        "    pytensor_mlx_mode = \"MLX\"\n",
        "    MLX_AVAILABLE = True\n",
        "except ImportError:\n",
        "    MLX_AVAILABLE = False\n",
        "\n",
        "def timer_jax(func, N=1000):\n",
        "    \"\"\"Time function execution with proper JAX synchronization, repeated N times\"\"\"\n",
        "    def wrapper(*args, **kwargs):\n",
        "        times = []\n",
        "        for _ in range(N):\n",
        "            start = time.perf_counter()\n",
        "            result = func(*args, **kwargs)\n",
        "            if hasattr(result, 'block_until_ready'):\n",
        "                result.block_until_ready()\n",
        "            elif isinstance(result, (list, tuple)):\n",
        "                for r in result:\n",
        "                    if hasattr(r, 'block_until_ready'):\n",
        "                        r.block_until_ready()\n",
        "            end = time.perf_counter()\n",
        "            times.append(end - start)\n",
        "        \n",
        "        mean_time = np.mean(times)\n",
        "        std_time = np.std(times)\n",
        "        return result, mean_time, std_time\n",
        "    return wrapper\n",
        "\n",
        "def timer_mlx(func, N=1000):\n",
        "    \"\"\"Time function execution with proper MLX synchronization, repeated N times\"\"\"\n",
        "    def wrapper(*args, **kwargs):\n",
        "        times = []\n",
        "        for _ in range(N):\n",
        "            start = time.perf_counter()\n",
        "            result = func(*args, **kwargs)\n",
        "            # For MLX, we need to use mx.eval() to force computation\n",
        "            if MLX_AVAILABLE:\n",
        "                if isinstance(result, (list, tuple)):\n",
        "                    mx.eval(*result)\n",
        "                else:\n",
        "                    mx.eval(result)\n",
        "            end = time.perf_counter()\n",
        "            times.append(end - start)\n",
        "        \n",
        "        mean_time = np.mean(times)\n",
        "        std_time = np.std(times)\n",
        "        return result, mean_time, std_time\n",
        "    return wrapper\n",
        "\n",
        "def run_benchmark(N=1000):\n",
        "    \"\"\"Run comprehensive benchmark comparing PyTensor JAX vs MLX backends\"\"\"\n",
        "    import pandas as pd\n",
        "    \n",
        "    sizes = [2, 4, 1080, 2080, 3080]\n",
        "    results = []\n",
        "    \n",
        "    print(f\"Running benchmarks with N={N} repetitions per test...\")\n",
        "    \n",
        "    for size in sizes:\n",
        "        print(f\"Testing {size}x{size} matrices...\")\n",
        "        \n",
        "        # Generate test matrices with fixed seed for reproducibility\n",
        "        np.random.seed(42)\n",
        "        A = np.random.randn(size, size).astype(np.float32)\n",
        "        B = np.random.randn(size, size).astype(np.float32)\n",
        "        C = np.random.randn(size, size).astype(np.float32)\n",
        "\n",
        "        pt_A = pt.matrix('A', dtype='float32')\n",
        "        pt_B = pt.matrix('B', dtype='float32')  \n",
        "        pt_C = pt.matrix('C', dtype='float32')\n",
        "        result = pt.dot(pt.dot(pt_A, pt_B), pt_C)\n",
        "\n",
        "\n",
        "        f_jax = function([pt_A, pt_B, pt_C], result, mode=pytensor_jax_mode, trust_input=True)\n",
        "        f_mlx = function([pt_A, pt_B, pt_C], result, mode=pytensor_mlx_mode, trust_input=True)\n",
        "        f_jax(A, B, C)\n",
        "        f_mlx(A, B, C)\n",
        "        \n",
        "        # === TEST 1: Matrix Multiplication Chain ===\n",
        "        # PyTensor + JAX backend\n",
        "        @timer_jax\n",
        "        def pytensor_jax_matmul():\n",
        "            return f_jax(A, B, C)\n",
        "        \n",
        "        # PyTensor + MLX backend\n",
        "        @timer_mlx\n",
        "        def pytensor_mlx_matmul():\n",
        "            if not MLX_AVAILABLE:\n",
        "                return None, float('inf'), 0\n",
        "            return f_mlx(A, B, C)\n",
        "        \n",
        "        # Run matrix multiplication test\n",
        "        _, jax_mean, jax_std = pytensor_jax_matmul()\n",
        "        try:\n",
        "            _, mlx_mean, mlx_std = pytensor_mlx_matmul()\n",
        "        except Exception as e:\n",
        "            print(f\"MLX matmul error: {e}\")\n",
        "            mlx_mean, mlx_std = float('inf'), 0\n",
        "        \n",
        "        # Calculate percentage improvement (positive = MLX is faster, negative = MLX is slower)\n",
        "        if mlx_mean != float('inf') and mlx_mean > 0:\n",
        "            speedup_percentage = ((jax_mean - mlx_mean) / jax_mean) * 100\n",
        "            speedup_str = f'{speedup_percentage:+.1f}%'\n",
        "        else:\n",
        "            speedup_str = 'N/A'\n",
        "        \n",
        "        results.append({\n",
        "            'Size': f'{size}x{size}',\n",
        "            'Operation': 'Matrix Chain (A @ B @ C)',\n",
        "            'PyTensor+JAX Mean (s)': f'{jax_mean:.6f}',\n",
        "            'PyTensor+JAX Std (s)': f'{jax_std:.6f}',\n",
        "            'PyTensor+MLX Mean (s)': f'{mlx_mean:.6f}' if mlx_mean != float('inf') else 'Error',\n",
        "            'PyTensor+MLX Std (s)': f'{mlx_std:.6f}' if mlx_mean != float('inf') else 'N/A',\n",
        "            'MLX Performance': speedup_str\n",
        "        })\n",
        "        \n",
        "        # === TEST 2: Element-wise Operations ===\n",
        "        # PyTensor + JAX\n",
        "        result = pt.sin(pt_A) + pt.cos(pt_B)\n",
        "        f_jax = function([pt_A, pt_B], result, mode=pytensor_jax_mode, trust_input=True)\n",
        "        f_mlx = function([pt_A, pt_B], result, mode=pytensor_mlx_mode, trust_input=True)\n",
        "        f_jax(A, B)\n",
        "        f_mlx(A, B)\n",
        "\n",
        "        @timer_jax\n",
        "        def pytensor_jax_elemwise():\n",
        "            return f_jax(A, B)\n",
        "        \n",
        "        # PyTensor + MLX\n",
        "        @timer_mlx\n",
        "        def pytensor_mlx_elemwise():\n",
        "            if not MLX_AVAILABLE:\n",
        "                return None, float('inf'), 0\n",
        "            return f_mlx(A, B)\n",
        "        \n",
        "        # Run element-wise test\n",
        "        _, jax_mean, jax_std = pytensor_jax_elemwise()\n",
        "        try:\n",
        "            _, mlx_mean, mlx_std = pytensor_mlx_elemwise()\n",
        "        except Exception as e:\n",
        "            print(f\"MLX elemwise error: {e}\")\n",
        "            mlx_mean, mlx_std = float('inf'), 0\n",
        "        \n",
        "        # Calculate percentage improvement\n",
        "        if mlx_mean != float('inf') and mlx_mean > 0:\n",
        "            speedup_percentage = ((jax_mean - mlx_mean) / jax_mean) * 100\n",
        "            speedup_str = f'{speedup_percentage:+.1f}%'\n",
        "        else:\n",
        "            speedup_str = 'N/A'\n",
        "        \n",
        "        results.append({\n",
        "            'Size': f'{size}x{size}',\n",
        "            'Operation': 'Element-wise (sin(A) + cos(B))',\n",
        "            'PyTensor+JAX Mean (s)': f'{jax_mean:.6f}',\n",
        "            'PyTensor+JAX Std (s)': f'{jax_std:.6f}',\n",
        "            'PyTensor+MLX Mean (s)': f'{mlx_mean:.6f}' if mlx_mean != float('inf') else 'Error',\n",
        "            'PyTensor+MLX Std (s)': f'{mlx_std:.6f}' if mlx_mean != float('inf') else 'N/A',\n",
        "            'MLX Performance': speedup_str\n",
        "        })\n",
        "        \n",
        "        # === TEST 3: Matrix Addition with Broadcasting ===\n",
        "        # PyTensor + JAX\n",
        "        result = pt_A + pt_B.T\n",
        "        f_jax = function([pt_A, pt_B], result, mode=pytensor_jax_mode, trust_input=True)\n",
        "        f_mlx = function([pt_A, pt_B], result, mode=pytensor_mlx_mode, trust_input=True)\n",
        "        f_jax(A, B)\n",
        "        f_mlx(A, B)\n",
        "        @timer_jax\n",
        "        def pytensor_jax_broadcast():\n",
        "            return f_jax(A, B)\n",
        "        \n",
        "        # PyTensor + MLX\n",
        "        @timer_mlx\n",
        "        def pytensor_mlx_broadcast():\n",
        "            if not MLX_AVAILABLE:\n",
        "                return None, float('inf'), 0\n",
        "            return f_mlx(A, B)\n",
        "        \n",
        "        # Run broadcasting test\n",
        "        _, jax_mean, jax_std = pytensor_jax_broadcast()\n",
        "        try:\n",
        "            _, mlx_mean, mlx_std = pytensor_mlx_broadcast()\n",
        "        except Exception as e:\n",
        "            print(f\"MLX broadcast error: {e}\")\n",
        "            mlx_mean, mlx_std = float('inf'), 0\n",
        "        \n",
        "        # Calculate percentage improvement\n",
        "        if mlx_mean != float('inf') and mlx_mean > 0:\n",
        "            speedup_percentage = ((jax_mean - mlx_mean) / jax_mean) * 100\n",
        "            speedup_str = f'{speedup_percentage:+.1f}%'\n",
        "        else:\n",
        "            speedup_str = 'N/A'\n",
        "        \n",
        "        results.append({\n",
        "            'Size': f'{size}x{size}',\n",
        "            'Operation': 'Broadcasting (A + B.T)',\n",
        "            'PyTensor+JAX Mean (s)': f'{jax_mean:.6f}',\n",
        "            'PyTensor+JAX Std (s)': f'{jax_std:.6f}',\n",
        "            'PyTensor+MLX Mean (s)': f'{mlx_mean:.6f}' if mlx_mean != float('inf') else 'Error',\n",
        "            'PyTensor+MLX Std (s)': f'{mlx_std:.6f}' if mlx_mean != float('inf') else 'N/A',\n",
        "            'MLX Performance': speedup_str\n",
        "        })\n",
        "    \n",
        "    # Create and display results table\n",
        "    df = pd.DataFrame(results)\n",
        "    return df\n",
        "\n",
        "def main(N=1000):\n",
        "    \"\"\"Main benchmark execution\"\"\"\n",
        "    # Display system info\n",
        "    system_info = {\n",
        "        'JAX version': jax.__version__,\n",
        "        'PyTensor version': pytensor.__version__,\n",
        "        'MLX Available': 'Yes' if MLX_AVAILABLE else 'No',\n",
        "        'Platform': 'Apple Silicon' if MLX_AVAILABLE else 'Generic',\n",
        "        'Repetitions (N)': N\n",
        "    }\n",
        "    \n",
        "    if MLX_AVAILABLE:\n",
        "        system_info['MLX version'] = mx.__version__\n",
        "    \n",
        "    import pandas as pd\n",
        "    info_df = pd.DataFrame([system_info])\n",
        "    \n",
        "    # Then run benchmarks\n",
        "    results_df = run_benchmark(N=N)\n",
        "    \n",
        "    return info_df, results_df\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Running benchmarks with N=150 repetitions per test...\n",
            "Testing 2x2 matrices...\n",
            "Testing 4x4 matrices...\n",
            "Testing 1080x1080 matrices...\n",
            "Testing 2080x2080 matrices...\n",
            "Testing 3080x3080 matrices...\n"
          ]
        }
      ],
      "source": [
        "iteration=150\n",
        "_, results = main(N=iteration)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Benchmark Results over 150 repetitions:\n",
            "     Size                      Operation PyTensor+JAX Mean (s) PyTensor+JAX Std (s) PyTensor+MLX Mean (s) PyTensor+MLX Std (s) MLX Performance\n",
            "      2x2       Matrix Chain (A @ B @ C)              0.000009             0.000002              0.000305             0.000299        -3213.5%\n",
            "      2x2 Element-wise (sin(A) + cos(B))              0.000007             0.000002              0.000352             0.003757        -5078.0%\n",
            "      2x2         Broadcasting (A + B.T)              0.000007             0.000001              0.000188             0.000153        -2721.1%\n",
            "      4x4       Matrix Chain (A @ B @ C)              0.000009             0.000001              0.000209             0.000063        -2126.2%\n",
            "      4x4 Element-wise (sin(A) + cos(B))              0.000007             0.000001              0.000180             0.000066        -2449.5%\n",
            "      4x4         Broadcasting (A + B.T)              0.000007             0.000003              0.000181             0.000065        -2564.1%\n",
            "1080x1080       Matrix Chain (A @ B @ C)              0.005951             0.000356              0.001355             0.000392          +77.2%\n",
            "1080x1080 Element-wise (sin(A) + cos(B))              0.002820             0.000107              0.000432             0.000207          +84.7%\n",
            "1080x1080         Broadcasting (A + B.T)              0.000212             0.000035              0.000428             0.000206         -102.0%\n",
            "2080x2080       Matrix Chain (A @ B @ C)              0.027609             0.001255              0.004550             0.002528          +83.5%\n",
            "2080x2080 Element-wise (sin(A) + cos(B))              0.010086             0.000417              0.001175             0.000350          +88.3%\n",
            "2080x2080         Broadcasting (A + B.T)              0.000856             0.000068              0.001124             0.000241          -31.2%\n",
            "3080x3080       Matrix Chain (A @ B @ C)              0.093115             0.003823              0.013649             0.000513          +85.3%\n",
            "3080x3080 Element-wise (sin(A) + cos(B))              0.022586             0.000756              0.001930             0.000287          +91.5%\n",
            "3080x3080         Broadcasting (A + B.T)              0.002580             0.000161              0.001937             0.000257          +24.9%\n"
          ]
        }
      ],
      "source": [
        "print(f\"\\nBenchmark Results over {iteration} repetitions:\")\n",
        "print(results.to_string(index=False))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# # Additional timing analysis - separate compilation vs execution time\n",
        "# if MLX_AVAILABLE:\n",
        "#     print(\"\\n=== Detailed MLX Timing Analysis ===\")\n",
        "    \n",
        "#     # Test with medium-sized matrix\n",
        "#     np.random.seed(42)\n",
        "#     A = np.random.randn(512, 512).astype(np.float32)\n",
        "#     B = np.random.randn(512, 512).astype(np.float32)\n",
        "#     C = np.random.randn(512, 512).astype(np.float32)\n",
        "    \n",
        "#     # Create PyTensor function (compilation time)\n",
        "#     start = time.perf_counter()\n",
        "#     pt_A = pt.matrix('A', dtype='float32')\n",
        "#     pt_B = pt.matrix('B', dtype='float32')\n",
        "#     pt_C = pt.matrix('C', dtype='float32')\n",
        "#     result_expr = pt_A @ pt_B @ pt_C\n",
        "#     f_mlx = function([pt_A, pt_B, pt_C], result_expr, mode=pytensor_mlx_mode)\n",
        "#     compilation_time = time.perf_counter() - start\n",
        "    \n",
        "#     # First execution (may include additional compilation/optimization)\n",
        "#     start = time.perf_counter()\n",
        "#     result = f_mlx(A, B, C)\n",
        "#     mx.eval(result)  # Force evaluation\n",
        "#     first_exec_time = time.perf_counter() - start\n",
        "    \n",
        "#     # Subsequent executions (should be faster)\n",
        "#     exec_times = []\n",
        "#     for _ in range(1000):\n",
        "#         start = time.perf_counter()\n",
        "#         result = f_mlx(A, B, C)\n",
        "#         mx.eval(result)\n",
        "#         exec_times.append(time.perf_counter() - start)\n",
        "    \n",
        "#     avg_exec_time = np.mean(exec_times)\n",
        "#     std_exec_time = np.std(exec_times)\n",
        "    \n",
        "#     print(f\"Compilation time: {compilation_time:.4f}s\")\n",
        "#     print(f\"First execution: {first_exec_time:.4f}s\")\n",
        "#     print(f\"Average execution (5 runs): {avg_exec_time:.4f}s ± {std_exec_time:.4f}s\")\n",
        "#     print(f\"Individual execution times: {[f'{t:.4f}' for t in exec_times]}\")\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": []
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "mlx_env",
      "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
}
