Runnable Examples & Workloads

Complete, copy-pasteable JAX models ready to compile and execute on Qualcomm Snapdragon hardware.

1. Dense Layer: MatMul + ReLU

A classic fully-connected layer running via backend="qnn":

import jax
import jax.numpy as jnp
import jax_qnn

@jax.jit
def dense_layer(x, w, b):
    return jax.nn.relu(jnp.matmul(x, w) + b)

x = jnp.ones((128, 512), dtype=jnp.float32)
w = jnp.ones((512, 512), dtype=jnp.float32)
b = jnp.zeros((512,), dtype=jnp.float32)

y = jax.jit(dense_layer, backend="qnn")(x, w, b)
print("Output:", y.shape)

2. 2D Convolution (CNN Vision Block)

Image convolution with 3x3 kernel and bias addition:

import jax
import jax.numpy as jnp
import jax.lax as lax
import jax_qnn

@jax.jit
def conv_block(x, kernel, bias):
    conv_out = lax.conv_general_dilated(
        lhs=x, rhs=kernel, window_strides=(1, 1),
        padding='SAME', dimension_numbers=('NHWC', 'HWIO', 'NHWC')
    )
    return jax.nn.relu(conv_out + bias)

x = jnp.ones((1, 64, 64, 32), dtype=jnp.float32)
kernel = jnp.ones((3, 3, 32, 64), dtype=jnp.float32) / 9.0
bias = jnp.zeros((64,), dtype=jnp.float32)

out = jax.jit(conv_block, backend="qnn")(x, kernel, bias)
print("Output Shape:", out.shape)

3. Transformer Multi-Head Self-Attention

Scaled dot-product attention block targeting Hexagon NPU Vector Cores:

import jax
import jax.numpy as jnp
import jax_qnn

@jax.jit
def self_attention(q, k, v):
    d_k = q.shape[-1]
    scores = jnp.matmul(q, jnp.swapaxes(k, -2, -1)) / jnp.sqrt(d_k)
    weights = jax.nn.softmax(scores, axis=-1)
    return jnp.matmul(weights, v)

q = jnp.ones((1, 4, 128, 64), dtype=jnp.float32)
k = jnp.ones((1, 4, 128, 64), dtype=jnp.float32)
v = jnp.ones((1, 4, 128, 64), dtype=jnp.float32)

out = jax.jit(self_attention, backend="qnn")(q, k, v)
print("Attention Out:", out.shape)

4. Live Qualcomm Hardware Validation (Hexagon NPU & Adreno GPU)

Verify direct hardware execution on Snapdragon® X Elite accelerators:

python tools/test_qualcomm_hardware.py

Live Hardware Output (Snapdragon® X Elite on Windows 11 ARM64):

Loading native library: C:\Users\rawat\JAX-QNN\build\pjrt_qnn.dll

=======================================================
 Testing Qualcomm Target: Qualcomm Hexagon NPU (HTP) (ID=1)
=======================================================
[JAX-QNN] QnnInterface successfully bound (providers: 1)
 <W> Initializing HtpProvider
[JAX-QNN] QnnBackend created successfully.
 <W> Logs will be sent to the system's default channel
 <E> DspTransport.openSession qnn_open failed, 0x80000406, prio 100
 <E> IDspTransport: Unable to load lib 0x80000406
 <E> DspTransport.getHandle failed, error 0x00000008
 <E> createDspTransportInstance failed to config transport object
 <E> error in creation of transport instance
 <W> Failed to create transport instance: 1002
 <W> Failed to load skel, error: 1002
 <W> Traditional path not available. Switching to user driver path
 <W> HTP user driver is loaded. Switched to user driver path
[JAX-QNN] QnnContext created successfully.
[JAX-QNN] Native QNN Hardware Backend Initialized: C:\Users\rawat\Downloads\v2.37.0.250724\qairt\2.37.0.250724\lib\aarch64-windows-msvc\QnnHtp.dll
[+] Init Result: 0 (Backend: C:\Users\rawat\Downloads\v2.37.0.250724\qairt\2.37.0.250724\lib\aarch64-windows-msvc\QnnHtp.dll)
[JAX-QNN] Parsing StableHLO (425 bytes)
[JAX-QNN] Parsed computation 'main': 2 inputs, 1 outputs, 3 ops
[JAX-QNN]   Input 0: [4, 4] f32
[JAX-QNN]   Input 1: [4, 4] f32
[JAX-QNN]   Op: stablehlo.dot_general (inputs: 2, outputs: 0)
[JAX-QNN]   Op: stablehlo.maximum (inputs: 2, outputs: 1)
[JAX-QNN]   Op: stablehlo.return (inputs: 1, outputs: 0)
[+] Graph Compiled Successfully (Handle: 0x243694f18c0)
[+] Inputs: 2, Outputs: 1
[+] Hardware Execution Status: 0
[+] Correctness Verified! Output[0, 0] = 24.0 (Expected: 24.0)

=======================================================
 Testing Qualcomm Target: Qualcomm Adreno GPU (OpenCL/Direct3D) (ID=2)
=======================================================
 <W> Logs will be sent to the system's default channel
 <W> m_CFBCallbackInfoObj is not initialized, return emptyList
 <W> Logs will be sent to the system's default channel
 <W> Logs will be sent to the system's default channel
[JAX-QNN] QnnInterface successfully bound (providers: 1)
[JAX-QNN] QnnBackend created successfully.
[JAX-QNN] QnnContext created successfully.
[JAX-QNN] Native QNN Hardware Backend Initialized: C:\Users\rawat\Downloads\v2.37.0.250724\qairt\2.37.0.250724\lib\aarch64-windows-msvc\QnnGpu.dll
[+] Init Result: 0 (Backend: C:\Users\rawat\Downloads\v2.37.0.250724\qairt\2.37.0.250724\lib\aarch64-windows-msvc\QnnGpu.dll)
[JAX-QNN] Parsing StableHLO (425 bytes)
[JAX-QNN] Parsed computation 'main': 2 inputs, 1 outputs, 3 ops
[JAX-QNN]   Input 0: [4, 4] f32
[JAX-QNN]   Input 1: [4, 4] f32
[JAX-QNN]   Op: stablehlo.dot_general (inputs: 2, outputs: 0)
[JAX-QNN]   Op: stablehlo.maximum (inputs: 2, outputs: 1)
[JAX-QNN]   Op: stablehlo.return (inputs: 1, outputs: 0)
[+] Graph Compiled Successfully (Handle: 0x243800dc6f0)
[+] Inputs: 2, Outputs: 1
[+] Hardware Execution Status: 0
[+] Correctness Verified! Output[0, 0] = 24.0 (Expected: 24.0)