Accelerate JAX on
Snapdragon NPU & GPU

import jax
import jax.numpy as jnp
import jax_qnn

# 1. Discover Qualcomm hardware accelerators
print("Available Devices:", jax.devices("qnn"))

# 2. Define JAX Model
@jax.jit
def dense_model(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)

# 3. Compile & Execute directly on Qualcomm Hexagon NPU
output = jax.jit(dense_model, backend="qnn")(x, w, b)
print("Output shape:", output.shape)