← jax4jChapter 7 · Execution backends and measurement7 min read

Start on CPU, select a GPU engine

Choose tensor placement explicitly, inspect available providers and distinguish native GPU engines from the browser WebGPU bridge.

Java notebookJava kernel · starts on first runWebGPU · not requested

Run cells from top to bottom. Typed declarations and imports carry into later cells. Shift + Enter runs a cell. Reloading or leaving this page discards edits and Java state.

Ready to start

Four numbers, two possible places to compute

Squaring [1, 2, 3, 4] and summing gives 30. The interesting question here is where those operations happen. Can you identify the device that owns the input, select another one deliberately, and explain what crosses the boundary? A GPU selection deserves more evidence than a reassuring variable name.

A backend implements tensor operations. A device identifies a placement, such as cpu:0 or wgpu:0. A device descriptor is a request, not proof that an adapter exists. Constructing Device.wgpu() neither installs a library nor makes the browser grant access to a GPU.

The same tensor API has three execution paths: the core CPU backend, an optional native GPU engine through Java FFM and drivers, and this notebook's wgpu adapter through the emulated JVM and browser WebGPU. Native wgpu and the browser adapter are different implementations.

Choose Run all. CPU cells work independently of GPU availability; optional GPU cells print a clear message when no usable browser adapter is available. The first dependency loads both bundled core and the notebook bridge before backend discovery. Each lesson has its own session; leaving or reloading loses edits and Java state. The browser JVM uses a custom JDK 27, while a normal local jax4j application can use Java 25.

1. Start with a known CPU placement

We set the default backend to CPU before creating tensors. The system property is jax4j.tensor.backend; a native process can also set JAX4J_BACKEND. An explicit system property takes precedence over the environment variable. Pinning an unavailable default backend raises an error rather than quietly choosing a different engine.

The notebook's browser adapter has lower selection priority than CPU anyway. Writing the property makes our intention visible and keeps later experiments from changing the placement of the reference calculation.

In [ ]:

1. Confirm the CPU calculation

Expect cpu:0 -> 30.0. The intermediate square and scalar result are both owned tensors, so both appear in the resource list. A reduction changes shape and value; it does not move a tensor to another device.

Try this: change the fourth input to 5. The sum becomes 39. Rerun only this cell: you have checked the arithmetic without involving device selection. Restore the fixture when comparing the reference answers below.

2. Ask which providers are usable

A provider on the classpath and a working device are different things. Tensors.backends() returns available providers, while Tensors.devices() returns an ordinal-zero descriptor for each available provider. This is not a complete inventory of every physical GPU or a promise about multi-GPU support.

Availability checks can initialize native runtimes or request a browser adapter. They are useful diagnostics, but running them before a cold-start measurement has already warmed part of the system. Here that initialization is intentional.

In [ ]:

2. Inspect available providers and placements

CPU should be present. A usable bridge also reports a wgpu provider and wgpu:0. The browser must support WebGPU in this context and expose a suitable adapter. A hardware-enabled Chrome configuration is encouraging, but the request from this page is what establishes actual availability.

The dependency directive must remain in the first cell. Backend providers are discovered once by the registry. If you initialized a core-only session and then added the bridge, use Restart and run the revised setup first. Installing code after discovery does not refresh that static registry.

3. Select a device independently of the default

arrayOn places newly supplied host values on the requested device. Its explicit shape must describe the supplied data. We start with CPU again so you can separate the placement API from optional hardware requirements.

The default selects a location for factories that do not receive one. An explicit device selects the provider for that request. Changing the default later does not migrate tensors that already exist.

In [ ]:

3. Place host values explicitly

The answer remains 30, with input shape [4] on cpu:0. For a native application, Device.cuda(), Device.rocm(), Device.vulkan() and Device.wgpu() request the corresponding engine's first device. Supplying a descriptor for an unavailable engine fails; it does not convert the operation to CPU work.

4. Run the same primitive operations on browser WebGPU

This cell requests wgpu:0 only after the availability check. It deliberately prints a skipped message if the adapter is unavailable. There is no CPU substitute hidden behind the GPU label. A later device loss can still cause an error even after a successful availability check.

The bridge dispatches real browser GPU kernels for these operations. The guest JVM sends commands to JavaScript through the VM bridge, and JavaScript submits them to WebGPU. That communication contributes overhead of its own.

In [ ]:

4. Select the browser GPU explicitly

Expect 30 again when the GPU cell runs. Four numbers are an arithmetic fixture, not a workload likely to repay GPU setup and communication costs. Use the Monte Carlo introduction for a larger resident GPU calculation. Do not turn this tiny example into a speed claim by timing a single call.

5. Follow a transfer out and back

Tensor.on returns a tensor on the requested device. In jax4j 0.2.1 a cross-device move downloads through host memory and uploads to the target backend. It is not a direct GPU-to-GPU transfer. The destination has its own lifetime; the source remains valid and still needs closing.

A CPU tensor holding 1,2,3,4 moves through host download and target upload to a wgpu tensor, then returns through host staging to a new CPU tensor. Three distinct owned tensor handles are closed. A move to the same device instead returns the original handle.
In [ ]:

5. Move explicitly and keep both owners visible

The returned values should be [1, 2, 3, 4]. Cross-device transfers can be useful at application boundaries, but repeating this round trip inside an inner loop often defeats the point of resident GPU data. Place inputs once, do substantial work, and retrieve only the result you need.

There is an ownership exception worth remembering: .on with the tensor's existing device returns the same Java object. It does not give you another independently owned copy. The next lesson makes this aliasing behavior explicit.

6. Confirm the selected placement without another transfer

Operations inherit their input's placement. Inspecting .device() reports that placement without downloading tensor values. .toFloatArray() does download them, so keep reporting choices distinct from computation.

In [ ]:

6. Check placement through an operation

This final cell openly chooses CPU if no GPU is available, and prints that choice before doing any work. Expect [12, 13]. Change the selection to Device.cpu() and rerun: the placement changes, while this exact arithmetic answer stays the same.

Taking the selection into a native application

A local application needs core plus the matching optional engine artifact, compatible native libraries and a working driver. All artifacts use group com.nablatensor; use matching 0.2.1 release or locally built artifacts. All of these library JARs were confirmed on Maven Central on 21 October 2026. Publication does not establish hardware compatibility or operation coverage.

Requested placementOptional artifact0.2.1 validation
Device.cuda()jax4j-engine-cudaBuilds; execution last validated in 0.2.0
Device.rocm()jax4j-engine-rocmImplemented; hardware unverified
Device.vulkan()jax4j-engine-vulkanHardware validated
Device.wgpu()jax4j-engine-wgpuTensor hardware validated; no scalar replay adapter

OpenCL has no engine in current jax4j. Vector API SIMD is an optional scalar replay engine in jax4j-simd, not a tensor placement; it requires --add-modules jdk.incubator.vector and execution is unverified for 0.2.1. ROCm also has a published implementation but no hardware validation. Neither native engine runs in this browser notebook. See the current engine support matrix before choosing a native backend.

The native wgpu engine uses wgpu-native through FFM. This notebook uses its own browser adapter under the same wgpu selector. Native engine capabilities therefore cannot be inferred from the browser bridge. In particular, this bridge supports the eager primitives shown here but rejects generated JIT and device-region kernels. The next lessons keep compiled examples on CPU.

For installation details, use the jax4j README and native wgpu guide. Continue with native access, synchronization and resources to make completion and ownership as explicit as placement.


Questions or corrections? open an issue