{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Tensor Operations in Burn\n",
    "\n",
    "This notebook demonstrates basic tensor operations in Burn, a deep learning framework written in Rust."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "// Dependency declarations for the notebook.\n",
    "// The syntax is similar to Cargo.toml. Just prefix with :dep\n",
    "\n",
    ":dep burn = {path = \"../../crates/burn\"}\n",
    ":dep burn-flex = {path = \"../../crates/burn-flex\"}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "// Import packages\n",
    "use burn::prelude::*;\n",
    "use burn_flex::Flex;\n",
    "\n",
    "// Type alias for the backend (using CPU/Flex)\n",
    "type B = Flex;"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Tensor Creation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Empty tensor shape: Shape { dims: [2, 3, 4] }\n",
      "Zeros tensor: Tensor {\n",
      "  data:\n",
      "[[0.0, 0.0, 0.0],\n",
      " [0.0, 0.0, 0.0],\n",
      " [0.0, 0.0, 0.0]],\n",
      "  shape:  [3, 3],\n",
      "  device:  Cpu,\n",
      "  backend:  \"flex\",\n",
      "  kind:  \"Float\",\n",
      "  dtype:  \"f32\",\n",
      "}\n",
      "Ones tensor: Tensor {\n",
      "  data:\n",
      "[[1.0, 1.0, 1.0, 1.0],\n",
      " [1.0, 1.0, 1.0, 1.0]],\n",
      "  shape:  [2, 4],\n",
      "  device:  Cpu,\n",
      "  backend:  \"flex\",\n",
      "  kind:  \"Float\",\n",
      "  dtype:  \"f32\",\n",
      "}\n",
      "Full tensor (7.0): Tensor {\n",
      "  data:\n",
      "[[7.0, 7.0, 7.0],\n",
      " [7.0, 7.0, 7.0]],\n",
      "  shape:  [2, 3],\n",
      "  device:  Cpu,\n",
      "  backend:  \"flex\",\n",
      "  kind:  \"Float\",\n",
      "  dtype:  \"f32\",\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "let device = Default::default();\n",
    "\n",
    "// Create an empty tensor (uninitialized values)\n",
    "let empty: Tensor<B, 3> = Tensor::empty([2, 3, 4], &device);\n",
    "println!(\"Empty tensor shape: {:?}\", empty.shape());\n",
    "\n",
    "// Create a tensor filled with zeros\n",
    "let zeros: Tensor<B, 2> = Tensor::zeros([3, 3], &device);\n",
    "println!(\"Zeros tensor: {}\", zeros);\n",
    "\n",
    "// Create a tensor filled with ones\n",
    "let ones: Tensor<B, 2> = Tensor::ones([2, 4], &device);\n",
    "println!(\"Ones tensor: {}\", ones);\n",
    "\n",
    "// Create a tensor filled with a specific value\n",
    "let full: Tensor<B, 2> = Tensor::full([2, 3], 7.0, &device);\n",
    "println!(\"Full tensor (7.0): {}\", full);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "From slice:\n",
      "Tensor {\n",
      "  data:\n",
      "[[1.0, 2.0, 3.0],\n",
      " [4.0, 5.0, 6.0]],\n",
      "  shape:  [2, 3],\n",
      "  device:  Cpu,\n",
      "  backend:  \"flex\",\n",
      "  kind:  \"Float\",\n",
      "  dtype:  \"f32\",\n",
      "}\n",
      "Random tensor: Tensor {\n",
      "  data:\n",
      "[0.32371014, 0.41100568, 0.94457513, 0.8408601, 0.42262083],\n",
      "  shape:  [5],\n",
      "  device:  Cpu,\n",
      "  backend:  \"flex\",\n",
      "  kind:  \"Float\",\n",
      "  dtype:  \"f32\",\n",
      "}\n",
      "Normal distribution: Tensor {\n",
      "  data:\n",
      "[-0.22402725, 1.8367178, -1.1049407, -0.6302627, 1.1106112],\n",
      "  shape:  [5],\n",
      "  device:  Cpu,\n",
      "  backend:  \"flex\",\n",
      "  kind:  \"Float\",\n",
      "  dtype:  \"f32\",\n",
      "}\n",
      "Uniform [0, 10): Tensor {\n",
      "  data:\n",
      "[8.110331, 7.335061, 9.858947, 6.0834813, 3.6619747],\n",
      "  shape:  [5],\n",
      "  device:  Cpu,\n",
      "  backend:  \"flex\",\n",
      "  kind:  \"Float\",\n",
      "  dtype:  \"f32\",\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "// Create a tensor from a slice of values\n",
    "let data = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0];\n",
    "let from_slice = Tensor::<B, 1>::from_floats(data, &device).reshape([2, 3]);\n",
    "println!(\"From slice:\\n{}\", from_slice);\n",
    "\n",
    "// Create a random tensor\n",
    "use burn::tensor::Distribution;\n",
    "let random: Tensor<B, 1> = Tensor::random([5], Distribution::Default, &device);\n",
    "println!(\"Random tensor: {}\", random);\n",
    "\n",
    "// Create a tensor with normal distribution\n",
    "let normal: Tensor<B, 1> = Tensor::random([5], Distribution::Normal(0.0, 1.0), &device);\n",
    "println!(\"Normal distribution: {}\", normal);\n",
    "\n",
    "// Create a tensor with uniform distribution in range [0, 10)\n",
    "let uniform: Tensor<B, 1> = Tensor::random([5], Distribution::Uniform(0.0, 10.0), &device);\n",
    "println!(\"Uniform [0, 10): {}\", uniform);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Shape Operations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Original (2x3):\n",
      "Tensor {\n",
      "  data:\n",
      "[[1.0, 2.0, 3.0],\n",
      " [4.0, 5.0, 6.0]],\n",
      "  shape:  [2, 3],\n",
      "  device:  Cpu,\n",
      "  backend:  \"flex\",\n",
      "  kind:  \"Float\",\n",
      "  dtype:  \"f32\",\n",
      "}\n",
      "Reshaped (1x2x3): Tensor {\n",
      "  data:\n",
      "[[[1.0, 2.0, 3.0],\n",
      "  [4.0, 5.0, 6.0]]],\n",
      "  shape:  [1, 2, 3],\n",
      "  device:  Cpu,\n",
      "  backend:  \"flex\",\n",
      "  kind:  \"Float\",\n",
      "  dtype:  \"f32\",\n",
      "}\n",
      "Flattened: Tensor {\n",
      "  data:\n",
      "[1.0, 2.0, 3.0, 4.0, 5.0, 6.0],\n",
      "  shape:  [6],\n",
      "  device:  Cpu,\n",
      "  backend:  \"flex\",\n",
      "  kind:  \"Float\",\n",
      "  dtype:  \"f32\",\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "// Reshape tensor - change the dimensions without changing the data\n",
    "let tensor = Tensor::<B, 1>::from_floats([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &device).reshape([2, 3]);\n",
    "println!(\"Original (2x3):\\n{}\", tensor);\n",
    "\n",
    "let reshaped: Tensor<B, 3> = tensor.clone().reshape([1, 2, 3]);\n",
    "println!(\"Reshaped (1x2x3): {}\", reshaped);\n",
    "\n",
    "// Flatten - reshape to 1D\n",
    "let flat: Tensor<B, 1> = tensor.flatten(0, 1);\n",
    "println!(\"Flattened: {}\", flat);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Original:\n",
      "Tensor {\n",
      "  data:\n",
      "[[1.0, 2.0],\n",
      " [3.0, 4.0]],\n",
      "  shape:  [2, 2],\n",
      "  device:  Cpu,\n",
      "  backend:  \"flex\",\n",
      "  kind:  \"Float\",\n",
      "  dtype:  \"f32\",\n",
      "}\n",
      "Transposed:\n",
      "Tensor {\n",
      "  data:\n",
      "[[1.0, 3.0],\n",
      " [2.0, 4.0]],\n",
      "  shape:  [2, 2],\n",
      "  device:  Cpu,\n",
      "  backend:  \"flex\",\n",
      "  kind:  \"Float\",\n",
      "  dtype:  \"f32\",\n",
      "}\n",
      "Using .t():\n",
      "Tensor {\n",
      "  data:\n",
      "[[1.0, 3.0],\n",
      " [2.0, 4.0]],\n",
      "  shape:  [2, 2],\n",
      "  device:  Cpu,\n",
      "  backend:  \"flex\",\n",
      "  kind:  \"Float\",\n",
      "  dtype:  \"f32\",\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "// Transpose - swap dimensions\n",
    "let tensor = Tensor::<B, 1>::from_floats([1.0, 2.0, 3.0, 4.0], &device).reshape([2, 2]);\n",
    "println!(\"Original:\\n{}\", tensor);\n",
    "\n",
    "let transposed = tensor.clone().transpose();\n",
    "println!(\"Transposed:\\n{}\", transposed);\n",
    "\n",
    "// Also .t() works for 2D tensors\n",
    "let t = tensor.t();\n",
    "println!(\"Using .t():\\n{}\", t);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Before squeeze [1,1,2]: shape = Shape { dims: [1, 1, 2] }\n",
      "After squeeze: shape = Shape { dims: [2] }\n",
      "Before unsqueeze [2,2]: shape = Shape { dims: [2, 2] }\n",
      "After unsqueeze: shape = Shape { dims: [1, 2, 2] }\n"
     ]
    }
   ],
   "source": [
    "// Squeeze - remove dimensions of size 1\n",
    "let tensor = Tensor::<B, 1>::from_floats([1.0, 2.0], &device).reshape([1, 1, 2]);\n",
    "println!(\"Before squeeze [1,1,2]: shape = {:?}\", tensor.shape());\n",
    "\n",
    "let squeezed = tensor.squeeze::<1>();\n",
    "println!(\"After squeeze: shape = {:?}\", squeezed.shape());\n",
    "\n",
    "// Unsqueeze - add a dimension of size 1 at specified position\n",
    "let tensor = Tensor::<B, 1>::from_floats([1.0, 2.0, 3.0, 4.0], &device).reshape([2, 2]);\n",
    "println!(\"Before unsqueeze [2,2]: shape = {:?}\", tensor.shape());\n",
    "\n",
    "let unsqueezed = tensor.unsqueeze::<3>();\n",
    "println!(\"After unsqueeze: shape = {:?}\", unsqueezed.shape());"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Indexing and Slicing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Original tensor:\n",
      "Tensor {\n",
      "  data:\n",
      "[[1.0, 2.0, 3.0, 4.0],\n",
      " [5.0, 6.0, 7.0, 8.0],\n",
      " [9.0, 10.0, 11.0, 12.0]],\n",
      "  shape:  [3, 4],\n",
      "  device:  Cpu,\n",
      "  backend:  \"flex\",\n",
      "  kind:  \"Float\",\n",
      "  dtype:  \"f32\",\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "// Create a tensor for indexing examples\n",
    "let tensor = Tensor::<B, 1>::from_floats(\n",
    "    [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0, 12.0],\n",
    "&device\n",
    ").reshape([3, 4]);\n",
    "println!(\"Original tensor:\\n{}\", tensor);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Sliced [1..3, 1..4]:\n",
      "Tensor {\n",
      "  data:\n",
      "[[6.0, 7.0, 8.0],\n",
      " [10.0, 11.0, 12.0]],\n",
      "  shape:  [2, 3],\n",
      "  device:  Cpu,\n",
      "  backend:  \"flex\",\n",
      "  kind:  \"Float\",\n",
      "  dtype:  \"f32\",\n",
      "}\n",
      "Row 1: Tensor {\n",
      "  data:\n",
      "[[5.0, 6.0, 7.0, 8.0]],\n",
      "  shape:  [1, 4],\n",
      "  device:  Cpu,\n",
      "  backend:  \"flex\",\n",
      "  kind:  \"Float\",\n",
      "  dtype:  \"f32\",\n",
      "}\n",
      "Column 2: Tensor {\n",
      "  data:\n",
      "[[3.0],\n",
      " [7.0],\n",
      " [11.0]],\n",
      "  shape:  [3, 1],\n",
      "  device:  Cpu,\n",
      "  backend:  \"flex\",\n",
      "  kind:  \"Float\",\n",
      "  dtype:  \"f32\",\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "// Slice tensor - select a portion using ranges\n",
    "// Get rows 1-2 (index 1 to end), columns 1-3 (index 1 to 3)\n",
    "let sliced = tensor.clone().slice([1..3, 1..4]);\n",
    "println!(\"Sliced [1..3, 1..4]:\\n{}\", sliced);\n",
    "\n",
    "// Get single row\n",
    "let row = tensor.clone().slice([1..2, 0..4]);\n",
    "println!(\"Row 1: {}\", row);\n",
    "\n",
    "// Get single column\n",
    "let col = tensor.slice([0..3, 2..3]);\n",
    "println!(\"Column 2: {}\", col);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Basic Math Operations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "a = Tensor {\n",
      "  data:\n",
      "[[1.0, 2.0],\n",
      " [3.0, 4.0]],\n",
      "  shape:  [2, 2],\n",
      "  device:  Cpu,\n",
      "  backend:  \"flex\",\n",
      "  kind:  \"Float\",\n",
      "  dtype:  \"f32\",\n",
      "}\n",
      "b = Tensor {\n",
      "  data:\n",
      "[[5.0, 6.0],\n",
      " [7.0, 8.0]],\n",
      "  shape:  [2, 2],\n",
      "  device:  Cpu,\n",
      "  backend:  \"flex\",\n",
      "  kind:  \"Float\",\n",
      "  dtype:  \"f32\",\n",
      "}\n",
      "a + b = Tensor {\n",
      "  data:\n",
      "[[6.0, 8.0],\n",
      " [10.0, 12.0]],\n",
      "  shape:  [2, 2],\n",
      "  device:  Cpu,\n",
      "  backend:  \"flex\",\n",
      "  kind:  \"Float\",\n",
      "  dtype:  \"f32\",\n",
      "}\n",
      "a - b = Tensor {\n",
      "  data:\n",
      "[[-4.0, -4.0],\n",
      " [-4.0, -4.0]],\n",
      "  shape:  [2, 2],\n",
      "  device:  Cpu,\n",
      "  backend:  \"flex\",\n",
      "  kind:  \"Float\",\n",
      "  dtype:  \"f32\",\n",
      "}\n",
      "a * b = Tensor {\n",
      "  data:\n",
      "[[5.0, 12.0],\n",
      " [21.0, 32.0]],\n",
      "  shape:  [2, 2],\n",
      "  device:  Cpu,\n",
      "  backend:  \"flex\",\n",
      "  kind:  \"Float\",\n",
      "  dtype:  \"f32\",\n",
      "}\n",
      "a / b = Tensor {\n",
      "  data:\n",
      "[[0.2, 0.33333334],\n",
      " [0.42857143, 0.5]],\n",
      "  shape:  [2, 2],\n",
      "  device:  Cpu,\n",
      "  backend:  \"flex\",\n",
      "  kind:  \"Float\",\n",
      "  dtype:  \"f32\",\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "let a = Tensor::<B, 1>::from_floats([1.0, 2.0, 3.0, 4.0], &device).reshape([2, 2]);\n",
    "let b = Tensor::<B, 1>::from_floats([5.0, 6.0, 7.0, 8.0], &device).reshape([2, 2]);\n",
    "\n",
    "println!(\"a = {}\", a);\n",
    "println!(\"b = {}\", b);\n",
    "\n",
    "// Addition\n",
    "let c = a.clone() + b.clone();\n",
    "println!(\"a + b = {}\", c);\n",
    "\n",
    "// Subtraction\n",
    "let c = a.clone() - b.clone();\n",
    "println!(\"a - b = {}\", c);\n",
    "\n",
    "// Multiplication (element-wise)\n",
    "let c = a.clone() * b.clone();\n",
    "println!(\"a * b = {}\", c);\n",
    "\n",
    "// Division (element-wise)\n",
    "let c = a.clone() / b.clone();\n",
    "println!(\"a / b = {}\", c);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "a = Tensor {\n",
      "  data:\n",
      "[[1.0, 2.0],\n",
      " [3.0, 4.0]],\n",
      "  shape:  [2, 2],\n",
      "  device:  Cpu,\n",
      "  backend:  \"flex\",\n",
      "  kind:  \"Float\",\n",
      "  dtype:  \"f32\",\n",
      "}\n",
      "a + 10 = Tensor {\n",
      "  data:\n",
      "[[11.0, 12.0],\n",
      " [13.0, 14.0]],\n",
      "  shape:  [2, 2],\n",
      "  device:  Cpu,\n",
      "  backend:  \"flex\",\n",
      "  kind:  \"Float\",\n",
      "  dtype:  \"f32\",\n",
      "}\n",
      "a * 2 = Tensor {\n",
      "  data:\n",
      "[[2.0, 4.0],\n",
      " [6.0, 8.0]],\n",
      "  shape:  [2, 2],\n",
      "  device:  Cpu,\n",
      "  backend:  \"flex\",\n",
      "  kind:  \"Float\",\n",
      "  dtype:  \"f32\",\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "// Scalar operations\n",
    "let a = Tensor::<B, 1>::from_floats([1.0, 2.0, 3.0, 4.0], &device).reshape([2, 2]);\n",
    "\n",
    "println!(\"a = {}\", a);\n",
    "\n",
    "// Add scalar\n",
    "let c = a.clone() + 10.0;\n",
    "println!(\"a + 10 = {}\", c);\n",
    "\n",
    "// Multiply scalar\n",
    "let c = a.clone() * 2.0;\n",
    "println!(\"a * 2 = {}\", c);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "a = Tensor {\n",
      "  data:\n",
      "[[1.0, 2.0],\n",
      " [3.0, 4.0]],\n",
      "  shape:  [2, 2],\n",
      "  device:  Cpu,\n",
      "  backend:  \"flex\",\n",
      "  kind:  \"Float\",\n",
      "  dtype:  \"f32\",\n",
      "}\n",
      "b = Tensor {\n",
      "  data:\n",
      "[[5.0, 6.0],\n",
      " [7.0, 8.0]],\n",
      "  shape:  [2, 2],\n",
      "  device:  Cpu,\n",
      "  backend:  \"flex\",\n",
      "  kind:  \"Float\",\n",
      "  dtype:  \"f32\",\n",
      "}\n",
      "a @ b (matmul) = Tensor {\n",
      "  data:\n",
      "[[19.0, 22.0],\n",
      " [43.0, 50.0]],\n",
      "  shape:  [2, 2],\n",
      "  device:  Cpu,\n",
      "  backend:  \"flex\",\n",
      "  kind:  \"Float\",\n",
      "  dtype:  \"f32\",\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "// Matrix multiplication\n",
    "let a = Tensor::<B, 1>::from_floats([1.0, 2.0, 3.0, 4.0], &device).reshape([2, 2]);\n",
    "let b = Tensor::<B, 1>::from_floats([5.0, 6.0, 7.0, 8.0], &device).reshape([2, 2]);\n",
    "\n",
    "println!(\"a = {}\", a);\n",
    "println!(\"b = {}\", b);\n",
    "\n",
    "let result = a.matmul(b);\n",
    "println!(\"a @ b (matmul) = {}\", result);\n",
    "\n",
    "// Verify (rows of a · columns of b): row1 [1,2] · col1 [5,7] = 1*5+2*7 = 19, row1 [1,2] · col2 [6,8] = 1*6+2*8 = 22\n",
    "//                                      row2 [3,4] · col1 [5,7] = 3*5+4*7 = 43, row2 [3,4] · col2 [6,8] = 3*6+4*8 = 50"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Element-wise Math Functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "a = Tensor {\n",
      "  data:\n",
      "[0.0, 1.0, 2.0],\n",
      "  shape:  [3],\n",
      "  device:  Cpu,\n",
      "  backend:  \"flex\",\n",
      "  kind:  \"Float\",\n",
      "  dtype:  \"f32\",\n",
      "}\n",
      "exp(a) = Tensor {\n",
      "  data:\n",
      "[1.0, 2.7182817, 7.389056],\n",
      "  shape:  [3],\n",
      "  device:  Cpu,\n",
      "  backend:  \"flex\",\n",
      "  kind:  \"Float\",\n",
      "  dtype:  \"f32\",\n",
      "}\n",
      "log(a + 1) = Tensor {\n",
      "  data:\n",
      "[0.0, 0.6931472, 1.0986123],\n",
      "  shape:  [3],\n",
      "  device:  Cpu,\n",
      "  backend:  \"flex\",\n",
      "  kind:  \"Float\",\n",
      "  dtype:  \"f32\",\n",
      "}\n",
      "a.powf(2) = Tensor {\n",
      "  data:\n",
      "[0.0, 1.0, 4.0],\n",
      "  shape:  [3],\n",
      "  device:  Cpu,\n",
      "  backend:  \"flex\",\n",
      "  kind:  \"Float\",\n",
      "  dtype:  \"f32\",\n",
      "}\n",
      "a.powf(0.5) = Tensor {\n",
      "  data:\n",
      "[0.0, 1.0, 1.4142135],\n",
      "  shape:  [3],\n",
      "  device:  Cpu,\n",
      "  backend:  \"flex\",\n",
      "  kind:  \"Float\",\n",
      "  dtype:  \"f32\",\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "let a: Tensor<B, 1> = Tensor::from_floats([0.0, 1.0, 2.0], &device);\n",
    "\n",
    "println!(\"a = {}\", a);\n",
    "\n",
    "// Exponential\n",
    "println!(\"exp(a) = {}\", a.clone().exp());\n",
    "\n",
    "// Natural logarithm\n",
    "println!(\"log(a + 1) = {}\", (a.clone() + 1.0).log());\n",
    "\n",
    "// Power\n",
    "println!(\"a.powf(2) = {}\", a.clone().powf_scalar(2.0));\n",
    "println!(\"a.powf(0.5) = {}\", a.clone().powf_scalar(0.5));"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "angles = Tensor {\n",
      "  data:\n",
      "[0.0, 0.7853982, 1.5707964],\n",
      "  shape:  [3],\n",
      "  device:  Cpu,\n",
      "  backend:  \"flex\",\n",
      "  kind:  \"Float\",\n",
      "  dtype:  \"f32\",\n",
      "}\n",
      "sin(angles) = Tensor {\n",
      "  data:\n",
      "[0.0, 0.70710677, 1.0],\n",
      "  shape:  [3],\n",
      "  device:  Cpu,\n",
      "  backend:  \"flex\",\n",
      "  kind:  \"Float\",\n",
      "  dtype:  \"f32\",\n",
      "}\n",
      "cos(angles) = Tensor {\n",
      "  data:\n",
      "[1.0, 0.70710677, -4.371139e-8],\n",
      "  shape:  [3],\n",
      "  device:  Cpu,\n",
      "  backend:  \"flex\",\n",
      "  kind:  \"Float\",\n",
      "  dtype:  \"f32\",\n",
      "}\n",
      "tan(angles) = Tensor {\n",
      "  data:\n",
      "[0.0, 1.0, -22877332.0],\n",
      "  shape:  [3],\n",
      "  device:  Cpu,\n",
      "  backend:  \"flex\",\n",
      "  kind:  \"Float\",\n",
      "  dtype:  \"f32\",\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "// Trigonometric functions\n",
    "let angles: Tensor<B, 1> = Tensor::from_floats([0.0, std::f32::consts::PI / 4.0, std::f32::consts::PI / 2.0], &device);\n",
    "\n",
    "println!(\"angles = {}\", angles);\n",
    "println!(\"sin(angles) = {}\", angles.clone().sin());\n",
    "println!(\"cos(angles) = {}\", angles.clone().cos());\n",
    "println!(\"tan(angles) = {}\", angles.clone().tan());"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6. Reduction Operations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Tensor:\n",
      "Tensor {\n",
      "  data:\n",
      "[[1.0, 2.0, 3.0],\n",
      " [4.0, 5.0, 6.0]],\n",
      "  shape:  [2, 3],\n",
      "  device:  Cpu,\n",
      "  backend:  \"flex\",\n",
      "  kind:  \"Float\",\n",
      "  dtype:  \"f32\",\n",
      "}\n",
      "Sum: Tensor {\n",
      "  data:\n",
      "[21.0],\n",
      "  shape:  [1],\n",
      "  device:  Cpu,\n",
      "  backend:  \"flex\",\n",
      "  kind:  \"Float\",\n",
      "  dtype:  \"f32\",\n",
      "}\n",
      "Mean: Tensor {\n",
      "  data:\n",
      "[3.5],\n",
      "  shape:  [1],\n",
      "  device:  Cpu,\n",
      "  backend:  \"flex\",\n",
      "  kind:  \"Float\",\n",
      "  dtype:  \"f32\",\n",
      "}\n",
      "Product: Tensor {\n",
      "  data:\n",
      "[720.0],\n",
      "  shape:  [1],\n",
      "  device:  Cpu,\n",
      "  backend:  \"flex\",\n",
      "  kind:  \"Float\",\n",
      "  dtype:  \"f32\",\n",
      "}\n",
      "Max: Tensor {\n",
      "  data:\n",
      "[6.0],\n",
      "  shape:  [1],\n",
      "  device:  Cpu,\n",
      "  backend:  \"flex\",\n",
      "  kind:  \"Float\",\n",
      "  dtype:  \"f32\",\n",
      "}\n",
      "Min: Tensor {\n",
      "  data:\n",
      "[1.0],\n",
      "  shape:  [1],\n",
      "  device:  Cpu,\n",
      "  backend:  \"flex\",\n",
      "  kind:  \"Float\",\n",
      "  dtype:  \"f32\",\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "let tensor = Tensor::<B, 1>::from_floats([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &device).reshape([2, 3]);\n",
    "println!(\"Tensor:\\n{}\", tensor);\n",
    "\n",
    "// Sum all elements\n",
    "println!(\"Sum: {}\", tensor.clone().sum());\n",
    "\n",
    "// Mean of all elements\n",
    "println!(\"Mean: {}\", tensor.clone().mean());\n",
    "\n",
    "// Product of all elements\n",
    "println!(\"Product: {}\", tensor.clone().prod());\n",
    "\n",
    "// Maximum and minimum\n",
    "println!(\"Max: {}\", tensor.clone().max());\n",
    "println!(\"Min: {}\", tensor.clone().min());"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Tensor:\n",
      "Tensor {\n",
      "  data:\n",
      "[[1.0, 2.0, 3.0],\n",
      " [4.0, 5.0, 6.0]],\n",
      "  shape:  [2, 3],\n",
      "  device:  Cpu,\n",
      "  backend:  \"flex\",\n",
      "  kind:  \"Float\",\n",
      "  dtype:  \"f32\",\n",
      "}\n",
      "Sum dim 0: Tensor {\n",
      "  data:\n",
      "[[5.0, 7.0, 9.0]],\n",
      "  shape:  [1, 3],\n",
      "  device:  Cpu,\n",
      "  backend:  \"flex\",\n",
      "  kind:  \"Float\",\n",
      "  dtype:  \"f32\",\n",
      "}\n",
      "Sum dim 1: Tensor {\n",
      "  data:\n",
      "[[6.0],\n",
      " [15.0]],\n",
      "  shape:  [2, 1],\n",
      "  device:  Cpu,\n",
      "  backend:  \"flex\",\n",
      "  kind:  \"Float\",\n",
      "  dtype:  \"f32\",\n",
      "}\n",
      "Mean dim 0: Tensor {\n",
      "  data:\n",
      "[[2.5, 3.5, 4.5]],\n",
      "  shape:  [1, 3],\n",
      "  device:  Cpu,\n",
      "  backend:  \"flex\",\n",
      "  kind:  \"Float\",\n",
      "  dtype:  \"f32\",\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "// Reduce along specific dimensions\n",
    "let tensor = Tensor::<B, 1>::from_floats([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &device).reshape([2, 3]);\n",
    "println!(\"Tensor:\\n{}\", tensor);\n",
    "\n",
    "// Sum along dimension 0 (columns)\n",
    "println!(\"Sum dim 0: {}\", tensor.clone().sum_dim(0));\n",
    "\n",
    "// Sum along dimension 1 (rows)\n",
    "println!(\"Sum dim 1: {}\", tensor.clone().sum_dim(1));\n",
    "\n",
    "// Mean along dimension 0\n",
    "println!(\"Mean dim 0: {}\", tensor.clone().mean_dim(0));"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 7. Comparison and Selection"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "a = Tensor {\n",
      "  data:\n",
      "[1.0, 5.0, 3.0, 8.0],\n",
      "  shape:  [4],\n",
      "  device:  Cpu,\n",
      "  backend:  \"flex\",\n",
      "  kind:  \"Float\",\n",
      "  dtype:  \"f32\",\n",
      "}\n",
      "b = Tensor {\n",
      "  data:\n",
      "[4.0, 2.0, 6.0, 7.0],\n",
      "  shape:  [4],\n",
      "  device:  Cpu,\n",
      "  backend:  \"flex\",\n",
      "  kind:  \"Float\",\n",
      "  dtype:  \"f32\",\n",
      "}\n",
      "a > b: Tensor {\n",
      "  data:\n",
      "[false, true, false, true],\n",
      "  shape:  [4],\n",
      "  device:  Cpu,\n",
      "  backend:  \"flex\",\n",
      "  kind:  \"Bool\",\n",
      "  dtype:  \"bool\",\n",
      "}\n",
      "a < b: Tensor {\n",
      "  data:\n",
      "[true, false, true, false],\n",
      "  shape:  [4],\n",
      "  device:  Cpu,\n",
      "  backend:  \"flex\",\n",
      "  kind:  \"Bool\",\n",
      "  dtype:  \"bool\",\n",
      "}\n",
      "a == b: Tensor {\n",
      "  data:\n",
      "[false, false, false, false],\n",
      "  shape:  [4],\n",
      "  device:  Cpu,\n",
      "  backend:  \"flex\",\n",
      "  kind:  \"Bool\",\n",
      "  dtype:  \"bool\",\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "let a: Tensor<B, 1> = Tensor::from_floats([1.0, 5.0, 3.0, 8.0], &device);\n",
    "let b: Tensor<B, 1> = Tensor::from_floats([4.0, 2.0, 6.0, 7.0], &device);\n",
    "\n",
    "println!(\"a = {}\", a);\n",
    "println!(\"b = {}\", b);\n",
    "\n",
    "// Element-wise comparison returns a boolean tensor\n",
    "let greater = a.clone().greater(b.clone());\n",
    "println!(\"a > b: {}\", greater);\n",
    "\n",
    "let less = a.clone().lower(b.clone());\n",
    "println!(\"a < b: {}\", less);\n",
    "\n",
    "let equal = a.clone().equal(b.clone());\n",
    "println!(\"a == b: {}\", equal);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Original: Tensor {\n",
      "  data:\n",
      "[1.0, 5.0, 3.0, 8.0],\n",
      "  shape:  [4],\n",
      "  device:  Cpu,\n",
      "  backend:  \"flex\",\n",
      "  kind:  \"Float\",\n",
      "  dtype:  \"f32\",\n",
      "}\n",
      "Where > 4, replace with 0: Tensor {\n",
      "  data:\n",
      "[1.0, 0.0, 3.0, 0.0],\n",
      "  shape:  [4],\n",
      "  device:  Cpu,\n",
      "  backend:  \"flex\",\n",
      "  kind:  \"Float\",\n",
      "  dtype:  \"f32\",\n",
      "}\n",
      "Where > 4, replace with -1: Tensor {\n",
      "  data:\n",
      "[1.0, -1.0, 3.0, -1.0],\n",
      "  shape:  [4],\n",
      "  device:  Cpu,\n",
      "  backend:  \"flex\",\n",
      "  kind:  \"Float\",\n",
      "  dtype:  \"f32\",\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "// Conditional selection\n",
    "let a: Tensor<B, 1> = Tensor::from_floats([1.0, 5.0, 3.0, 8.0], &device);\n",
    "\n",
    "// mask_where: where condition is true, use replacement value, else keep original value\n",
    "let condition = a.clone().greater_elem(4.0);\n",
    "let result = a.clone().mask_where(condition, Tensor::zeros([4], &device));\n",
    "println!(\"Original: {}\", a);\n",
    "println!(\"Where > 4, replace with 0: {}\", result);\n",
    "\n",
    "// mask_fill: simpler - just replace values matching condition\n",
    "let result = a.clone().mask_fill(a.clone().greater_elem(4.0), -1.0);\n",
    "println!(\"Where > 4, replace with -1: {}\", result);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Summary\n",
    "\n",
    "In this notebook, we covered:\n",
    "- **Tensor Creation**: empty, zeros, ones, full, from_floats, random\n",
    "- **Shape Operations**: reshape, transpose, flatten, squeeze, unsqueeze\n",
    "- **Indexing and Slicing**: slice operation with ranges\n",
    "- **Math Operations**: add, sub, mul, div, matmul\n",
    "- **Element-wise Functions**: exp, log, powf_scalar, sin, cos, tan\n",
    "- **Reduction Operations**: sum, mean, prod, max, min\n",
    "- **Comparison**: greater, lower, equal, mask_where, mask_fill\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Rust",
   "language": "rust",
   "name": "rust"
  },
  "language_info": {
   "codemirror_mode": "rust",
   "file_extension": ".rs",
   "mimetype": "text/rust",
   "name": "rust",
   "pygment_lexer": "rust",
   "version": ""
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
