Docs / jax4j-core / com.nablatensor.jax4j.tensor.spi

abstract class

GpuTensorBackend

Shared F32 tensor kernels and graph lowering. Concrete backends own their native runtime.

View source ↗Full module reference on javadoc.io ↗

Methods

protected GpuRuntime runtime()

String deviceName()

DeviceBuffer upload(float[] data, Shape shape, DType dtype, Device device)

DeviceBuffer randomUniform(long seed, long counter, Shape shape, Device device)

DeviceBuffer randomNormal(long seed, long counter, Shape shape, Device device)

DeviceBuffer full(float value, Shape shape, DType dtype, Device device)

float[] download(DeviceBuffer buffer)

DeviceBuffer binary(Op op, DeviceBuffer a, DeviceBuffer b)

DeviceBuffer scalar(Op op, DeviceBuffer a, double value)

DeviceBuffer unary(Op op, DeviceBuffer a)

DeviceBuffer fused(Expr expr, DeviceBuffer[] inputs)

Runs a whole elementwise Expr chain as one generated GLSL compute shader (one dispatch, one intermediate buffer), instead of the SPI default of one dispatch and one buffer per node. The shader is emitted with the same SSA post-order walk as the CUDA backend, compiled to SPIR-V once by libshaderc, and cached (as a registered pipeline) by its source.

DeviceBuffer transpose(DeviceBuffer a)

DeviceBuffer matmul(DeviceBuffer a, DeviceBuffer b)

DeviceBuffer reduceSum(DeviceBuffer a)

DeviceBuffer reduceMax(DeviceBuffer a)

DeviceBuffer sumAxis0(DeviceBuffer a)

DeviceBuffer sumAxis(DeviceBuffer a, int axis, boolean keepDims)

DeviceBuffer maxAxis(DeviceBuffer a, int axis, boolean keepDims)

DeviceBuffer argMaxAxis(DeviceBuffer a, int axis)

DeviceBuffer maxAxisBackward(DeviceBuffer upstream, DeviceBuffer input, int axis)

DeviceBuffer reluBackward(DeviceBuffer upstream, DeviceBuffer input)

DeviceBuffer broadcastTo(DeviceBuffer a, Shape target)

DeviceBuffer conv2d(DeviceBuffer x, DeviceBuffer w, ConvSpec spec)

DeviceBuffer conv2dGradInput(DeviceBuffer upstream, DeviceBuffer w, ConvSpec spec)

DeviceBuffer conv2dGradWeight(DeviceBuffer x, DeviceBuffer upstream, ConvSpec spec)

DeviceBuffer reshape(DeviceBuffer a, Shape target)

DeviceBuffer sliceAxis0(DeviceBuffer a, int index)

DeviceBuffer stackAxis0(DeviceBuffer[] inputs)

DeviceBuffer batchedMatmul(DeviceBuffer a, DeviceBuffer b)

DeviceBuffer reindex(DeviceBuffer input, Shape target, int[] indices)

DeviceBuffer scatterAdd(DeviceBuffer input, Shape target, int[] indices)

DeviceBuffer program(com.nablatensor.jax4j.tensor.DeviceProgram program, DeviceBuffer[] inputs)

void sync()

void release(DeviceBuffer buffer)