torchvulkan 0.1.0a0__tar.gz

torchvulkan-0.1.0a0/.github/actions/setup-toolchain/action.yml

@@ -0,0 +1,90 @@
+name: Setup torchvulkan toolchain
+description: >
+  Install the build/runtime toolchain for the current runner OS: the Slang
+  compiler, a Vulkan driver + SDK headers, and spirv-cross. Exports the needed
+  PATH / env entries so later steps in the job can build and run torchvulkan.
+
+inputs:
+  slang-version:
+    description: Slang release version to install.
+    default: "2026.8.1"
+  vulkan-sdk-version:
+    description: Vulkan SDK version (Windows only).
+    default: "1.3.296.0"
+
+runs:
+  using: composite
+  steps:
+    - name: Resolve Slang arch
+      shell: bash
+      run: |
+        case "${{ runner.os }}-${{ runner.arch }}" in
+          Linux-X64)    echo "SLANG_ARCH=linux-x86_64"   >> "$GITHUB_ENV" ;;
+          macOS-ARM64)  echo "SLANG_ARCH=macos-aarch64"  >> "$GITHUB_ENV" ;;
+          macOS-X64)    echo "SLANG_ARCH=macos-x86_64"   >> "$GITHUB_ENV" ;;
+          Windows-X64)  echo "SLANG_ARCH=windows-x86_64" >> "$GITHUB_ENV" ;;
+          *) echo "::error::Unsupported runner ${{ runner.os }}-${{ runner.arch }}"; exit 1 ;;
+        esac
+
+    - name: Install Slang (macOS & Linux)
+      if: runner.os != 'Windows'
+      shell: bash
+      run: |
+        curl -L -o slang.zip "https://github.com/shader-slang/slang/releases/download/v${{ inputs.slang-version }}/slang-${{ inputs.slang-version }}-${SLANG_ARCH}.zip"
+        unzip slang.zip -d "$RUNNER_TEMP/slang"
+        echo "$RUNNER_TEMP/slang/bin" >> "$GITHUB_PATH"
+        echo "SLANG_DIR=$RUNNER_TEMP/slang/bin" >> "$GITHUB_ENV"
+
+    - name: Install Slang (Windows)
+      if: runner.os == 'Windows'
+      shell: pwsh
+      run: |
+        $zip = "slang-${{ inputs.slang-version }}-$env:SLANG_ARCH.zip"
+        Invoke-WebRequest -Uri "https://github.com/shader-slang/slang/releases/download/v${{ inputs.slang-version }}/$zip" -OutFile $zip
+        Expand-Archive -Path $zip -DestinationPath C:\slang
+        "C:\slang\bin" | Out-File -FilePath $env:GITHUB_PATH -Append
+        "SLANG_DIR=C:\slang\bin" | Out-File -FilePath $env:GITHUB_ENV -Append
+
+    - name: Install Vulkan (Linux)
+      if: runner.os == 'Linux'
+      shell: bash
+      run: |
+        wget -qO- https://packages.lunarg.com/lunarg-signing-key-pub.asc | sudo tee /etc/apt/trusted.gpg.d/lunarg.asc
+        sudo wget -qO /etc/apt/sources.list.d/lunarg-vulkan-$(lsb_release -cs).list http://packages.lunarg.com/vulkan/lunarg-vulkan-$(lsb_release -cs).list
+        sudo apt-get update
+        sudo apt-get install -y vulkan-sdk mesa-vulkan-drivers
+        vulkaninfo --summary
+
+    - name: Install Vulkan (macOS)
+      if: runner.os == 'macOS'
+      shell: bash
+      run: |
+        brew install molten-vk vulkan-tools ninja libxshmfence spirv-cross
+        MVK_ICD=$(find $(brew --prefix)/share $(brew --prefix)/etc $(brew --prefix molten-vk) -name "MoltenVK_icd.json" 2>/dev/null | head -n 1)
+        echo "VK_ICD_FILENAMES=$MVK_ICD" >> "$GITHUB_ENV"
+        echo "MVK_CONFIG_USE_METAL_ARGUMENT_BUFFERS=0" >> "$GITHUB_ENV"
+        echo "DYLD_LIBRARY_PATH=$(brew --prefix)/lib:$DYLD_LIBRARY_PATH" >> "$GITHUB_ENV"
+        export VK_ICD_FILENAMES="$MVK_ICD"
+        export MVK_CONFIG_USE_METAL_ARGUMENT_BUFFERS=0
+        export DYLD_LIBRARY_PATH="$(brew --prefix)/lib:$DYLD_LIBRARY_PATH"
+        vulkaninfo --summary
+
+    - name: Install Vulkan (Windows)
+      if: runner.os == 'Windows'
+      shell: pwsh
+      run: |
+        Invoke-WebRequest -Uri "https://github.com/pal1000/mesa-dist-win/releases/download/23.3.2/mesa3d-23.3.2-release-msvc.7z" -OutFile "mesa.7z"
+        7z x mesa.7z -omesa -y | Out-Null
+
+        $lvpJson = (Resolve-Path "mesa\x64\lvp_icd.x86_64.json").Path
+        $driverKey = "HKLM:\SOFTWARE\Khronos\Vulkan\Drivers"
+        if (!(Test-Path $driverKey)) { New-Item -Path $driverKey -Force | Out-Null }
+        New-ItemProperty -Path $driverKey -Name $lvpJson -PropertyType DWord -Value 0 -Force | Out-Null
+
+        $sdkVersion = "${{ inputs.vulkan-sdk-version }}"
+        $installer = "VulkanSDK-$sdkVersion-Installer.exe"
+        Invoke-WebRequest -Uri "https://sdk.lunarg.com/sdk/download/$sdkVersion/windows/$installer" -OutFile $installer
+        Start-Process -FilePath ".\$installer" -ArgumentList "--accept-licenses --default-answer --confirm-command install" -Wait -NoNewWindow
+
+        "VULKAN_SDK=C:\VulkanSDK\$sdkVersion" | Out-File -FilePath $env:GITHUB_ENV -Append
+        "C:\VulkanSDK\$sdkVersion\Bin" | Out-File -FilePath $env:GITHUB_PATH -Append

torchvulkan-0.1.0a0/.github/scripts/install-toolchain-linux.sh

@@ -0,0 +1,27 @@
+#!/usr/bin/env bash
+# Installs the torchvulkan build toolchain inside the cibuildwheel manylinux
+# container (AlmaLinux 8 / manylinux_2_28): the Vulkan SDK (headers + loader +
+# the spirv-cross executable) and the Slang compiler (slangc).
+#
+# The install paths here MUST match the VULKAN_SDK / SLANG_DIR / PATH entries in
+# [tool.cibuildwheel.linux].environment in pyproject.toml.
+set -euo pipefail
+
+SLANG_VERSION="2026.8.1"
+VULKAN_SDK_VERSION="1.3.296.0"
+
+dnf install -y wget tar xz unzip
+
+wget -q -O /tmp/vulkansdk.tar.xz \
+  "https://sdk.lunarg.com/sdk/download/${VULKAN_SDK_VERSION}/linux/vulkansdk-linux-x86_64-${VULKAN_SDK_VERSION}.tar.xz"
+mkdir -p /opt/vulkan
+tar -xf /tmp/vulkansdk.tar.xz -C /opt/vulkan --strip-components=1
+
+wget -q -O /tmp/slang.zip \
+  "https://github.com/shader-slang/slang/releases/download/v${SLANG_VERSION}/slang-${SLANG_VERSION}-linux-x86_64.zip"
+mkdir -p /opt/slang
+unzip -q /tmp/slang.zip -d /opt/slang
+
+test -f /opt/vulkan/x86_64/include/vulkan/vulkan.h
+test -x /opt/vulkan/x86_64/bin/spirv-cross
+test -x /opt/slang/bin/slangc

torchvulkan-0.1.0a0/.github/workflows/release.yaml

@@ -0,0 +1,99 @@
+name: Release
+
+on:
+  push:
+    tags:
+      - "v*"
+  workflow_dispatch: {}
+
+permissions:
+  contents: read
+
+jobs:
+  build_sdist:
+    name: Build sdist
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          submodules: recursive
+          fetch-depth: 0
+
+      - uses: actions/setup-python@v5
+        with:
+          python-version: "3.10"
+
+      - name: Build sdist
+        run: |
+          python -m pip install --upgrade pip build
+          python -m build --sdist
+
+      - uses: actions/upload-artifact@v4
+        with:
+          name: dist-sdist
+          path: dist/*.tar.gz
+
+  build_wheels:
+    name: Wheels (${{ matrix.name }})
+    runs-on: ${{ matrix.os }}
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+          - name: Linux x86_64
+            os: ubuntu-latest
+          - name: macOS Apple Silicon
+            os: macos-15
+          - name: Windows x86_64
+            os: windows-latest
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          submodules: recursive
+          fetch-depth: 0
+
+      - uses: actions/setup-python@v5
+        with:
+          python-version: "3.12"
+
+      - name: Compute version for setuptools_scm
+        shell: bash
+        run: |
+          python -m pip install --upgrade setuptools_scm
+          echo "SETUPTOOLS_SCM_PRETEND_VERSION_FOR_TORCHVULKAN=$(python -m setuptools_scm)" >> "$GITHUB_ENV"
+
+      - name: Install toolchain (Slang, Vulkan, spirv-cross)
+        if: runner.os != 'Linux'
+        uses: ./.github/actions/setup-toolchain
+
+      - name: Build wheels
+        uses: pypa/cibuildwheel@v3.2.1
+        env:
+          CIBW_BUILD_VERBOSITY: "1"
+
+      - uses: actions/upload-artifact@v4
+        with:
+          name: dist-${{ matrix.os }}
+          path: wheelhouse/*.whl
+
+  pypi-publish:
+    name: Upload release to PyPI
+    needs: [build_sdist, build_wheels]
+    if: startsWith(github.ref, 'refs/tags/')
+    runs-on: ubuntu-latest
+    environment:
+      name: pypi
+      url: https://pypi.org/p/torchvulkan
+    permissions:
+      id-token: write
+    steps:
+      - uses: actions/download-artifact@v4
+        with:
+          path: dist
+          pattern: dist-*
+          merge-multiple: true
+
+      - uses: pypa/gh-action-pypi-publish@release/v1
+        with:
+          packages-dir: dist
+          skip-existing: true

torchvulkan-0.1.0a0/.github/workflows/test_torchvulkan.yml

@@ -0,0 +1,60 @@
+name: torchvulkan CI
+
+on:
+  push:
+    branches: [ main ]
+  pull_request:
+    branches: [ main ]
+
+jobs:
+  build-and-test:
+    name: ${{ matrix.os }} / Py${{ matrix.python-version }} / ${{ matrix.torch-factor }}
+    runs-on: ${{ matrix.os }}
+    env:
+      SKBUILD_CMAKE_GENERATOR: Ninja
+    strategy:
+      fail-fast: false
+      matrix:
+        os: [ubuntu-latest, macos-15, windows-latest]
+        python-version: ["3.10", "3.12", "3.14"]
+        torch-factor: ["torch210", "torch213"]
+
+    steps:
+      - name: Checkout Code
+        uses: actions/checkout@v4
+        with:
+          submodules: recursive
+
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python-version }}
+
+      - name: Install toolchain (Slang, Vulkan, spirv-cross)
+        uses: ./.github/actions/setup-toolchain
+
+      - name: Install Tox
+        run: |
+          python -m pip install --upgrade pip
+          pip install tox
+
+      - name: Run Tests via Tox
+        env:
+          PIP_EXTRA_INDEX_URL: https://download.pytorch.org/whl/cpu
+          SKBUILD_CMAKE_GENERATOR: Ninja
+          MVK_CONFIG_LOG_LEVEL: 3
+          VK_LOADER_DEBUG: all
+        run: |
+          tox -e py-${{ matrix.torch-factor }}
+
+  testing-success:
+    name: testing success
+    needs: build-and-test
+    if: always()
+    runs-on: ubuntu-latest
+    steps:
+      - name: Verify all test jobs passed
+        if: needs.build-and-test.result != 'success'
+        run: |
+          echo "build-and-test result: ${{ needs.build-and-test.result }}"
+          exit 1

torchvulkan-0.1.0a0/.gitignore

@@ -0,0 +1,42 @@
+# torchvulkan
+torchvulkan/_version.py
+src/shaders/shader_registry.h
+
+# testing data
+data/
+
+# python
+.venv/
+venv/
+__pycache__/
+*.py[cod]
+*.pyo
+*.pyd
+.pytest_cache/
+.mypy_cache/
+
+# build files
+build/
+dist/
+*.egg-info/
+*.so
+*.dll
+*.dylib
+*.o
+*.a
+*.out
+
+# cmake
+CMakeCache.txt
+CMakeFiles/
+cmake_install.cmake
+Makefile
+
+# IDE files
+.vscode/
+.idea/
+*.swp
+
+# operating system
+.DS_Store
+Thumbs.db

torchvulkan-0.1.0a0/CMakeLists.txt

@@ -0,0 +1,28 @@
+cmake_minimum_required(VERSION 3.18)
+project(torchvulkan LANGUAGES CXX)
+
+find_package(Torch REQUIRED)
+find_package(Vulkan REQUIRED)
+find_package(Python3 COMPONENTS Interpreter Development.Module REQUIRED)
+find_program(SLANGC_EXE slangc HINTS "$ENV{SLANG_DIR}" REQUIRED)
+find_program(SPIRV_CROSS_EXE spirv-cross REQUIRED)
+
+set(CMAKE_CXX_STANDARD 20)
+set(CMAKE_CXX_STANDARD_REQUIRED ON)
+
+add_subdirectory(src/shaders)
+
+file(GLOB_RECURSE SOURCES CONFIGURE_DEPENDS "src/*.cpp")
+
+python3_add_library(torchvulkan_lib MODULE ${SOURCES} WITH_SOABI)
+add_dependencies(torchvulkan_lib compile_shaders)
+
+set_target_properties(torchvulkan_lib PROPERTIES OUTPUT_NAME "_C" LIBRARY_OUTPUT_DIRECTORY "${CMAKE_CURRENT_SOURCE_DIR}/torchvulkan")
+target_link_libraries(torchvulkan_lib PRIVATE ${TORCH_LIBRARIES} Python3::Module ${CMAKE_DL_LIBS})
+target_include_directories(torchvulkan_lib PRIVATE src/ extern/vma extern/volk ${Vulkan_INCLUDE_DIRS})
+
+target_compile_definitions(torchvulkan_lib PRIVATE VK_NO_PROTOTYPES TORCH_EXTENSION_NAME=_C)
+
+target_compile_features(torchvulkan_lib PRIVATE cxx_std_20)
+
+install(TARGETS torchvulkan_lib DESTINATION torchvulkan)

torchvulkan-0.1.0a0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 torchvulkan
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

torchvulkan-0.1.0a0/PKG-INFO

@@ -0,0 +1,8 @@
+Metadata-Version: 2.2
+Name: torchvulkan
+Version: 0.1.0a0
+Requires-Dist: torch==2.13.0
+Provides-Extra: test
+Requires-Dist: pytest; extra == "test"
+Requires-Dist: tox; extra == "test"
+

torchvulkan-0.1.0a0/README.md

@@ -0,0 +1,163 @@
+<h1 align="center">torchvulkan</h1>
+<p align="center">
+  <i>A Cross-Platform Vulkan Backend for PyTorch</i>
+</p>
+
+<p align="center">
+  <a href="https://github.com/zainsharief/torchvulkan/actions/workflows/test_torchvulkan.yml"><img alt="CI" src="https://github.com/zainsharief/torchvulkan/actions/workflows/test_torchvulkan.yml/badge.svg"></a>
+  <a href="LICENSE"><img alt="License: MIT" src="https://img.shields.io/badge/License-MIT-blue.svg"></a>
+  <img alt="Python 3.10+" src="https://img.shields.io/badge/python-3.10%2B-blue">
+  <img alt="Status: experimental" src="https://img.shields.io/badge/status-experimental-orange">
+</p>
+
+`torchvulkan` registers a `vulkan` device with PyTorch and runs tensor operations as [Vulkan](https://www.vulkan.org/) compute shaders. Vulkan runs on GPUs from AMD, Intel, NVIDIA, and Apple (via MoltenVK), plus software rasterizers, so you get one vendor-neutral GPU backend across Linux, macOS, and Windows.
+
+You use it like any other PyTorch device:
+
+```python
+import torch
+import torchvulkan  # registers the 'vulkan' device on import
+
+x = torch.randn(1024, 1024).to("vulkan")
+y = torch.randn(1024, 1024).to("vulkan")
+
+z = (x @ y).relu()      # runs on the GPU via Vulkan compute shaders
+print(z.sum().item())   # move a scalar back to the CPU
+```
+
+> [!WARNING]
+> **Experimental.** `torchvulkan` is an early-stage project. A small set of operators is implemented natively; everything else falls back to the CPU. Expect breaking changes before `1.0`.
+
+## Table of contents
+
+- [Installation](#installation)
+- [Quickstart](#quickstart)
+- [Example: training MNIST](#example-training-mnist)
+- [Python API](#python-api)
+- [CPU fallback](#cpu-fallback)
+- [Limitations](#limitations)
+- [How it works](#how-it-works)
+
+## Installation
+
+```bash
+pip install torchvulkan
+```
+
+At runtime you need a **Vulkan driver**: your GPU's driver on Linux and Windows (or a software rasterizer such as Mesa [lavapipe](https://docs.mesa3d.org/drivers/llvmpipe.html)), or [MoltenVK](https://github.com/KhronosGroup/MoltenVK) on macOS. Verify it worked:
+
+```python
+import torchvulkan
+print(torchvulkan.is_available(), torchvulkan.device_count())
+```
+
+<details>
+<summary><b>Building from source</b></summary>
+
+Prebuilt wheels are published for Linux (x86_64), macOS (Apple Silicon), and Windows (x86_64) on Python 3.10 to 3.14. On other platforms `pip` builds from the source distribution, which needs the shader toolchain below. The same steps build a local checkout for development.
+
+| Requirement | Notes |
+| --- | --- |
+| A Vulkan driver / loader | Your GPU's Vulkan driver, or a software rasterizer such as Mesa [lavapipe](https://docs.mesa3d.org/drivers/llvmpipe.html). On macOS this is [MoltenVK](https://github.com/KhronosGroup/MoltenVK). |
+| [Vulkan SDK](https://vulkan.lunarg.com/) | Provides Vulkan headers and `spirv-cross`. |
+| [`slangc`](https://github.com/shader-slang/slang/releases) | The Slang shader compiler, on your `PATH`. |
+| `spirv-cross` | Ships with the Vulkan SDK (or your package manager). |
+| CMake ≥ 3.18 + a C++20 compiler | Ninja is recommended (`SKBUILD_CMAKE_GENERATOR=Ninja`). |
+| `torch==2.13.0` | Install before building; the extension is compiled against this exact version. |
+
+> The exact per-OS commands CI uses to install this toolchain live in [`.github/actions/setup-toolchain`](.github/actions/setup-toolchain/action.yml).
+
+```bash
+pip install torch==2.13.0
+pip install .            # editable install for development: pip install -e .
+```
+
+</details>
+
+## Quickstart
+
+```python
+import torch
+import torchvulkan
+
+# Move tensors to the Vulkan device
+a = torch.tensor([[1.0, 2.0], [3.0, 4.0]], device="vulkan")
+b = torch.ones(2, 2, device="vulkan")
+
+c = a @ b + 1.0          # matmul + broadcasted add
+d = c.exp().sum(dim=0)   # elementwise + reduction
+
+print(d.cpu())           # bring the result back to the CPU
+```
+
+`nn.Module`s move to the device the usual way:
+
+```python
+model = torch.nn.Linear(128, 10).to("vulkan")
+out = model(torch.randn(32, 128).to("vulkan"))
+```
+
+## Example: training MNIST
+
+[`examples/mnist.py`](examples/mnist.py) trains a small fully-connected network on the Vulkan device end to end. Forward, loss, backward, and an Adam step all run through the backend:
+
+```bash
+pip install torchvision
+python examples/mnist.py
+```
+
+```python
+device = torch.device("vulkan" if torchvulkan.is_available() else "cpu")
+model = SimpleNN().to(device)
+criterion = torch.nn.CrossEntropyLoss()
+optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
+
+for data, targets in train_loader:
+    data, targets = data.to(device), targets.to(device)
+    loss = criterion(model(data), targets)
+    optimizer.zero_grad()
+    loss.backward()
+    optimizer.step()
+```
+
+## Python API
+
+```python
+import torchvulkan
+
+torchvulkan.is_available()   # -> bool: a usable Vulkan device was found
+torchvulkan.device_count()   # -> int: number of Vulkan devices
+torchvulkan.synchronize()    # block until queued GPU work has finished
+torchvulkan.empty_cache()    # release cached device memory back to the driver
+torchvulkan.__version__      # package version
+```
+
+Importing `torchvulkan` registers the backend, after which the device is addressed as `"vulkan"` / `torch.device("vulkan")` throughout PyTorch.
+
+## CPU fallback
+
+By default, any operation, dtype, or tensor shape the backend doesn't implement natively runs on the **CPU** (with a one-time warning), so models keep working despite partial coverage.
+
+Set `TORCHVULKAN_STRICT=1` to turn those silent fallbacks into errors. This is useful for finding which operators a workload actually needs:
+
+```bash
+TORCHVULKAN_STRICT=1 python your_script.py
+```
+
+## Limitations
+
+- `bfloat16` and complex dtypes not natively supported; falls back to the CPU.
+- Tensor rank ≤ 4. Tensors with more than 4 dimensions fall back to the CPU.
+- `torch>=2.10.0` only. The extension is compiled against 2.13.0 but has testing support for >=2.10.0.
+- APIs and coverage may change before `1.0.0`.
+
+## How it works
+
+1. Importing the package registers a PyTorch `PrivateUse1` backend and renames it to `vulkan`.
+2. Operations dispatched to a `vulkan` tensor are routed to native implementations that record and submit **Vulkan compute shaders**. The shaders are written in [Slang](https://github.com/shader-slang/slang), compiled to SPIR-V at build time, and embedded into the extension.
+3. Device memory is managed with [VMA](https://github.com/GPUOpen-LibrariesAndSDKs/VulkanMemoryAllocator); Vulkan entry points are loaded with [volk](https://github.com/zeux/volk).
+4. Anything without a native implementation is handled by [CPU fallback](#cpu-fallback).
+
+
+> [!NOTE]
+> **Disclaimer:** `torchvulkan` is an independent open-source project and is not affiliated with or endorsed by PyTorch or the Linux Foundation.

torchvulkan-0.1.0a0/examples/main.py

@@ -0,0 +1,10 @@
+import torch
+import torchvulkan as torchvk
+
+if __name__ == "__main__":
+    try:
+        assert torchvk.is_available()
+        print(f"torchvulkan is available! Found {torchvk.device_count()} device(s).")
+
+    except Exception as e:
+        print(f"An error occurred while checking torchvulkan availability: {e}")

torchvulkan-0.1.0a0/examples/mnist.py

@@ -0,0 +1,51 @@
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torchvision import datasets, transforms
+from torch.utils.data import DataLoader
+import torchvulkan as torchvk
+
+batch_size = 32
+learning_rate = 0.001
+epochs = 3
+device = torch.device('vulkan' if torchvk.is_available() else 'cpu')
+
+transform = transforms.ToTensor()
+train_dataset = datasets.MNIST(root='./data', train=True, transform=transform, download=True)
+train_loader = DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True)
+
+class SimpleNN(nn.Module):
+    def __init__(self):
+        super(SimpleNN, self).__init__()
+        self.flatten = nn.Flatten()
+        self.fc1 = nn.Linear(28 * 28, 128)
+        self.relu = nn.ReLU()
+        self.fc2 = nn.Linear(128, 10)
+
+    def forward(self, x):
+        x = self.flatten(x)
+        x = self.fc1(x)
+        x = self.relu(x)
+        x = self.fc2(x)
+        return x
+
+model = SimpleNN().to(device)
+
+criterion = nn.CrossEntropyLoss()
+optimizer = optim.Adam(model.parameters(), lr=learning_rate)
+
+for epoch in range(epochs):
+    model.train()
+    for batch_idx, (data, targets) in enumerate(train_loader):
+        data, targets = data.to(device), targets.to(device)
+
+        scores = model(data)
+        loss = criterion(scores, targets)
+
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+
+        print(f"\repoch [{epoch+1}/{epochs}] | batch {batch_idx+1}/{len(train_loader)} | loss: {loss.item():.4f}", end='')
+
+print("Training complete!")