← jax4jChapter 0 · Getting started5 min read

What belongs in jax4j?

Explore the boundary between jax4j and Quant, then run a tensor computation in a Java notebook.

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

Square three numbers and add them up. A finance library and a neural network can both need that operation; neither needs to own its implementation. That shared numerical foundation is where jax4j starts.

jax4j owns tensors, automatic differentiation, JIT compilation, batching, linear algebra and neural network utilities. Quant builds financial models on that foundation: products, Monte Carlo valuation, calibration, risk and CVA. The Java packages make the boundary visible: compute types live under com.nablatensor.jax4j, while the finance APIs live under com.nablatensor.quant and related finance packages.

Try the boundary in a notebook

Run the cells in order, or choose Run all. The first run starts a Java session in your browser and downloads the runtime pieces it needs. It can take a little while; the status line shows progress. Your code runs locally in the browser; the page downloads runtime assets and requested library dependencies.

The first cell prints the guest Java version and declares a shape. The second uses that shape from the same session. The third loads the bundled jax4j 0.2.1 JAR and performs the actual tensor calculation. Its expected output is [1.0, 4.0, 9.0] and a sum of 14.0: each entry of [1, 2, 3] is multiplied by itself, then the three results are added.

In [ ]:

1. Meet the Java kernel

What carries between cells?

The notebook uses a persistent Java kernel. Explicitly typed declarations such as int elements = 6; become available to later cells, along with imports and method declarations. A local variable declared with var remains local to that cell. You can also enter a single expression to display its value. This notebook wraps cells into Java classes; it is not JShell.

Use Stop to terminate a running kernel, including a cell that does not finish. Restart kernel creates a fresh session and clears outputs while preserving your edited source. Run from the top again to restore declarations. Leaving or reloading the page discards both your edits and the session.

In [ ]:

2. Reuse state from the previous cell

In [ ]:

3. Run a real jax4j tensor computation

Try changing the third cell's inputs to {2, 3, 4}. Before you run it, work out the sum of squares. You should get 29. The try block closes each tensor when the computation finishes; later lessons will look more closely at tensor lifetimes.

From this browser to your application

The notebook compiles cells with TeaVM javac and executes their bytecode in a custom OpenJDK 27 Zero VM inside cxemu. The example library itself requires Java 25. This emulated runtime is useful for exploring small computations; its timings do not describe jax4j performance on your machine.

The third cell explicitly selects the CPU backend. The fourth selects Device.wgpu() and sends the tensor operations to Chrome's WebGPU device. Java still runs inside the VM; the WGSL compute shaders run through the browser. The adapter line reports the selected hardware. A software fallback is rejected rather than presented as GPU acceleration.

The notebook:jax4j-webgpu:0.2.1 dependency is a bundled notebook adapter, including jax4j-core 0.2.1. It is not a published Maven coordinate. Load it before the first tensor in a session so Java can discover both backends. CPU cells do not require WebGPU; if Chrome cannot provide a hardware adapter, the GPU cells report an error.

In [ ]:

4. Run tensors on browser WebGPU

The fifth cell multiplies a 2×3 matrix by a 3×2 matrix on both devices. Both results should be [58.0, 64.0, 139.0, 154.0]. This checks actual GPU dispatch and readback against the CPU implementation.

In [ ]:

5. Check GPU matrix multiplication against CPU

The sixth cell estimates π with Monte Carlo sampling. It generates 65,536 uniform (X, Y) pairs in the unit square. Points satisfying X² + Y² < 1 lie inside a quarter circle, so 4 × hits / samples estimates π. Increase samples to reduce sampling error; the result is an estimate, not an exact value.

Each row of xy contains a uniform (X, Y) pair generated on GPU. The fluent chain squares the coordinates, sums each row, constructs an inside/outside mask, counts the hits and scales by 4 / samples. Only .item() reads a scalar back to Java; all intermediate tensors stay on GPU.

BrowserWebGpuBackend.scope() is a notebook helper that releases the chain's temporary GPU buffers when the block ends. Keep tensors created in the scope inside that block. The displayed elapsed time includes random generation, GPU operations and scalar readback; the first run also includes setup and shader compilation. Chaining uses separate dispatches, not fused JIT compilation.

In [ ]:

6. Monte Carlo π on GPU

The browser bridge currently supports F32 arithmetic, unary operations, matrix multiplication, reductions, transpose, fill, copy, broadcasting and seeded uniform draws. Generated JIT/GLSL kernels, normal random draws, convolutions and device control-flow regions are not supported yet. Stop and Restart close both the Java kernel and its GPU resources. Native CUDA and Vulkan are not exposed in this notebook. In your own Java application, optional engine dependencies provide those native execution paths. Tensor gradients, compilation and batching belong to the same compute project; option payoffs, market curves and capital calculations belong to Quant.

Continue with your first tensor on Java 25 for another hands-on notebook. For current setup instructions, use the jax4j quickstart, or explore the Quant curriculum for financial applications.


Questions or corrections? open an issue