Package io.github.kirstenali.deepj.tensor.metal

Class MetalBackend

java.lang.Object

io.github.kirstenali.deepj.tensor.metal.MetalBackend

All Implemented Interfaces:

CausalAttentionBackend, TensorBackend

public final class MetalBackend
extends Object

Constructor Summary

Constructors

Constructor	Description
MetalBackend()

Method Summary

Modifier and Type	Method	Description
void	adamWUpdate(Tensor w, Tensor g, Tensor mt, Tensor vt, float lr, float beta1, float beta2, float eps, float weightDecay, float bc1, float bc2)
Tensor	add(Tensor a, Tensor b)
Tensor	addBroadcastCols(Tensor a, Tensor v)
Tensor	addBroadcastRows(Tensor a, Tensor v)
void	addInPlace(Tensor a, Tensor b)
Tensor	addRowVector(Tensor a, Tensor v)
Tensor	addScalar(Tensor a, float scalar)
void	addScalarInPlace(Tensor a, float s)
Tensor	batchedMatmul(Tensor left, Tensor right, int batches, boolean transposeLeft, boolean transposeRight)
Tensor	cast(Tensor input, DType dtype)
CausalAttentionResult	causalAttention(Tensor query, Tensor key, Tensor value, int heads, float scale)
CausalAttentionResult	causalAttention(Tensor query, Tensor key, Tensor value, int queryHeads, int keyValueHeads, float scale)
AttentionGradients	causalAttentionBackward(Tensor query, Tensor key, Tensor value, CausalAttentionResult saved, Tensor gradient, int heads, float scale)
AttentionGradients	causalAttentionBackward(Tensor query, Tensor key, Tensor value, CausalAttentionResult saved, Tensor gradient, int queryHeads, int keyValueHeads, float scale)
Tensor	causalMask(Tensor input, int sequenceLength)
Tensor	causalSoftmax(Tensor input, int sequenceLength, float scale)
Tensor	clamp(Tensor a, float min, float max)
Tensor	copy(Tensor source)	Independent storage with the same float bits, without recording an autograd operation.
void	copyInto(Tensor target, Tensor source)	Explicitly assigns logical values, casting to the destination's storage dtype.
CrossEntropyResult	crossEntropy(Tensor logits, int[] targets)
CrossEntropyResult	crossEntropy(Tensor logits, int[] targets, boolean[] mask)
Tensor	crossEntropyGradient(Tensor logits, int[] targets)
Tensor	crossEntropyGradient(Tensor logits, int[] targets, boolean[] mask)
float	crossEntropyLoss(Tensor logits, int[] targets)
float	crossEntropyLoss(Tensor logits, int[] targets, boolean[] mask)
Tensor	divide(Tensor a, Tensor b)
Tensor	divideBroadcastCols(Tensor a, Tensor v)
void	divideInPlace(Tensor a, Tensor b)
Tensor	divideScalar(Tensor a, float scalar)
void	divideScalarInPlace(Tensor a, float s)
Tensor	exp(Tensor a)
void	expInPlace(Tensor a)
Tensor	gelu(Tensor a)
Tensor	geluBackward(Tensor input, Tensor gradOutput)
void	geluInPlace(Tensor a)
Tensor	headPrefixMatmul(Tensor left, Tensor cache, int heads, int length, boolean transposeCache)
void	indexAdd(Tensor target, int axis, int[] indices, Tensor updates)
Tensor	indexSelect(Tensor input, int axis, int[] indices)
Tensor	indexSelectGradient(Tensor gradient, int[] shape, int axis, int[] indices)
static boolean	isAvailable()
float	l2Norm(List<Tensor> tensors)
Tensor	layerNormBackward(Tensor dXHat, Tensor xHat, Tensor std, int dim)
Tensor	log(Tensor a)
void	logInPlace(Tensor a)
void	materializeTensor(Tensor t)
Tensor	matmul(Tensor a, Tensor b)
Tensor	maxAlongRows(Tensor a)
Tensor	maximumAxes(Tensor input, int[] axes, boolean keep)
Tensor	maximumAxesBackward(Tensor input, Tensor maximum, Tensor gradient, int[] axes, boolean keep)
Tensor	meanAlongRows(Tensor a)
Tensor	meanAxes(Tensor input, int[] axes, boolean keep)
Tensor	mergeHeads(Tensor input, int heads)
Tensor	multiply(Tensor a, Tensor b)
Tensor	multiplyBroadcastCols(Tensor a, Tensor v)
Tensor	multiplyBroadcastRows(Tensor a, Tensor v)
void	multiplyInPlace(Tensor a, Tensor b)
Tensor	multiplyScalar(Tensor a, float scalar)
void	multiplyScalarInPlace(Tensor a, float s)
Tensor	neg(Tensor a)
void	negInPlace(Tensor a)
Tensor	normalizationBackward(Tensor gradient, Tensor normalized, Tensor denominator, Tensor gamma, boolean centered)
NormalizationResult	normalize(Tensor input, Tensor gamma, Tensor beta, float epsilon, boolean centered)
Tensor	pow(Tensor a, float exponent)
void	releaseResources()
void	releaseTemporaryResources()
Tensor	relu(Tensor a)
Tensor	reluBackward(Tensor input, Tensor gradOutput)
void	reluInPlace(Tensor a)
RmsNormResult	rmsNorm(Tensor input, Tensor gamma, float epsilon)
Tensor	rmsNormBackward(Tensor gradient, Tensor normalized, Tensor rms, Tensor gamma)
Tensor	rotary(Tensor input, Tensor cosine, Tensor sine, int sequenceLength, boolean inverse)
void	scatterAddRows(Tensor target, int[] indices, Tensor grad)
Tensor	sigmoid(Tensor a)
void	sigmoidInPlace(Tensor a)
Tensor	sliceGradient(Tensor gradient, int[] shape, int axis, int start, int step)
Tensor	sliceRows(Tensor input, int[] rows)
Tensor	softmaxAxis(Tensor input, int axis)
Tensor	softmaxBackward(Tensor gradOutput, Tensor softmaxOut)
Tensor	softmaxBackward(Tensor gradient, Tensor probabilities, float scale)
Tensor	softmaxRows(Tensor logits)
Tensor	splitHeads(Tensor input, int heads)
Tensor	sqrt(Tensor a)
void	sqrtInPlace(Tensor a)
Tensor	subtract(Tensor a, Tensor b)
Tensor	subtractBroadcastCols(Tensor a, Tensor v)
void	subtractInPlace(Tensor a, Tensor b)
float	sum(Tensor a)
float	sumAbs(Tensor a)
Tensor	sumAlongCols(Tensor a)
Tensor	sumAlongRows(Tensor a)
Tensor	sumAxes(Tensor input, int[] axes, boolean keep)
Tensor	sumRows(Tensor a)
boolean	supportsCausalAttention(Tensor query, Tensor key, Tensor value, int heads, float scale)
boolean	supportsCausalAttention(Tensor query, Tensor key, Tensor value, int queryHeads, int keyValueHeads, float scale)
static boolean	supportsNative16Matmul(DType dtype)	Whether compact matrix products can use native16 operands with FP32 accumulation.
Tensor	swiGlu(Tensor gate, Tensor up)
SwiGluBackwardResult	swiGluBackward(Tensor gradient, Tensor gate, Tensor up)
Tensor	tanh(Tensor a)
void	tanhInPlace(Tensor a)
Tensor	transpose(Tensor a)
boolean	usesDeviceStorage()	Whether backend copies reside on a device rather than in ordinary host storage.
Tensor	varianceAlongRows(Tensor a)
Tensor	varianceAxes(Tensor input, int[] axes, boolean keep, int correction)
void	writeHeadCache(Tensor cache, Tensor update, int heads, int offset)
Tensor	zeros(int rows, int columns)	Backend-owned positive-zero storage; ordinary Tensor constructors stay CPU-backed.
Tensor	zeros(DType dtype, int... shape)

Methods inherited from class java.lang.Object

clone, equals, finalize, getClass, hashCode, notify, notifyAll, toString, wait, wait, wait

Methods inherited from interface io.github.kirstenali.deepj.tensor.TensorBackend

adamWUpdate, add, addBroadcastCols, addBroadcastRows, addInPlace, addRowVector, addScalar, addScalarInPlace, batchedMatmul, cast, causalMask, causalSoftmax, clamp, copy, copyInto, crossEntropy, crossEntropy, crossEntropyGradient, crossEntropyGradient, crossEntropyLoss, crossEntropyLoss, divide, divideBroadcastCols, divideInPlace, divideScalar, divideScalarInPlace, exp, expand, expInPlace, gelu, geluBackward, geluInPlace, headPrefixMatmul, indexAdd, indexSelect, indexSelectGradient, l2Norm, layerNormBackward, log, logInPlace, materializeTensor, matmul, maxAlongRows, maximumAxes, maximumAxesBackward, meanAlongRows, meanAxes, mergeHeads, multiply, multiplyBroadcastCols, multiplyBroadcastRows, multiplyInPlace, multiplyScalar, multiplyScalarInPlace, neg, negInPlace, normalizationBackward, normalize, permute, pow, releaseResources, releaseTemporaryResources, relu, reluBackward, reluInPlace, reshape, rmsNorm, rmsNormBackward, rotary, scatterAddRows, select, sigmoid, sigmoidInPlace, slice, sliceGradient, sliceRows, softmaxAxis, softmaxAxisBackward, softmaxBackward, softmaxBackward, softmaxRows, splitHeads, sqrt, sqrtInPlace, subtract, subtractBroadcastCols, subtractInPlace, sum, sumAbs, sumAlongCols, sumAlongRows, sumAxes, sumRows, sumToShape, swiGlu, swiGluBackward, tanh, tanhInPlace, transpose, unsqueeze, usesDeviceStorage, varianceAlongRows, varianceAxes, view, writeHeadCache, zeros, zeros

Constructor Details

MetalBackend

public MetalBackend()

Method Details

supportsCausalAttention

public boolean supportsCausalAttention(Tensor query,
 Tensor key,
 Tensor value,
 int heads,
 float scale)

Specified by:

supportsCausalAttention in interface CausalAttentionBackend

supportsCausalAttention

public boolean supportsCausalAttention(Tensor query,
 Tensor key,
 Tensor value,
 int queryHeads,
 int keyValueHeads,
 float scale)

Specified by:

supportsCausalAttention in interface CausalAttentionBackend

causalAttention

public CausalAttentionResult causalAttention(Tensor query,
 Tensor key,
 Tensor value,
 int heads,
 float scale)

Specified by:

causalAttention in interface CausalAttentionBackend

causalAttention

public CausalAttentionResult causalAttention(Tensor query,
 Tensor key,
 Tensor value,
 int queryHeads,
 int keyValueHeads,
 float scale)

Specified by:

causalAttention in interface CausalAttentionBackend

causalAttentionBackward

public AttentionGradients causalAttentionBackward(Tensor query,
 Tensor key,
 Tensor value,
 CausalAttentionResult saved,
 Tensor gradient,
 int heads,
 float scale)

Specified by:

causalAttentionBackward in interface CausalAttentionBackend

causalAttentionBackward

public AttentionGradients causalAttentionBackward(Tensor query,
 Tensor key,
 Tensor value,
 CausalAttentionResult saved,
 Tensor gradient,
 int queryHeads,
 int keyValueHeads,
 float scale)

Specified by:

causalAttentionBackward in interface CausalAttentionBackend

writeHeadCache

public void writeHeadCache(Tensor cache,
 Tensor update,
 int heads,
 int offset)

headPrefixMatmul

public Tensor headPrefixMatmul(Tensor left,
 Tensor cache,
 int heads,
 int length,
 boolean transposeCache)

adamWUpdate

public void adamWUpdate(Tensor w,
 Tensor g,
 Tensor mt,
 Tensor vt,
 float lr,
 float beta1,
 float beta2,
 float eps,
 float weightDecay,
 float bc1,
 float bc2)

swiGlu

public Tensor swiGlu(Tensor gate,
 Tensor up)

swiGluBackward

public SwiGluBackwardResult swiGluBackward(Tensor gradient,
 Tensor gate,
 Tensor up)

scatterAddRows

public void scatterAddRows(Tensor target,
 int[] indices,
 Tensor grad)

releaseTemporaryResources

public void releaseTemporaryResources()

releaseResources

public void releaseResources()

crossEntropyGradient

public Tensor crossEntropyGradient(Tensor logits,
 int[] targets)

crossEntropy

public CrossEntropyResult crossEntropy(Tensor logits,
 int[] targets)

crossEntropyLoss

public float crossEntropyLoss(Tensor logits,
 int[] targets)

crossEntropyLoss

public float crossEntropyLoss(Tensor logits,
 int[] targets,
 boolean[] mask)

crossEntropy

public CrossEntropyResult crossEntropy(Tensor logits,
 int[] targets,
 boolean[] mask)

crossEntropyGradient

public Tensor crossEntropyGradient(Tensor logits,
 int[] targets,
 boolean[] mask)

layerNormBackward

public Tensor layerNormBackward(Tensor dXHat,
 Tensor xHat,
 Tensor std,
 int dim)

rmsNorm

public RmsNormResult rmsNorm(Tensor input,
 Tensor gamma,
 float epsilon)

rmsNormBackward

public Tensor rmsNormBackward(Tensor gradient,
 Tensor normalized,
 Tensor rms,
 Tensor gamma)

normalize

public NormalizationResult normalize(Tensor input,
 Tensor gamma,
 Tensor beta,
 float epsilon,
 boolean centered)

normalizationBackward

public Tensor normalizationBackward(Tensor gradient,
 Tensor normalized,
 Tensor denominator,
 Tensor gamma,
 boolean centered)

addInPlace

public void addInPlace(Tensor a,
 Tensor b)

subtractInPlace

public void subtractInPlace(Tensor a,
 Tensor b)

multiplyInPlace

public void multiplyInPlace(Tensor a,
 Tensor b)

divideInPlace

public void divideInPlace(Tensor a,
 Tensor b)

multiplyScalarInPlace

public void multiplyScalarInPlace(Tensor a,
 float s)

addScalarInPlace

public void addScalarInPlace(Tensor a,
 float s)

divideScalarInPlace

public void divideScalarInPlace(Tensor a,
 float s)

sqrtInPlace

public void sqrtInPlace(Tensor a)

negInPlace

public void negInPlace(Tensor a)

expInPlace

public void expInPlace(Tensor a)

logInPlace

public void logInPlace(Tensor a)

reluInPlace

public void reluInPlace(Tensor a)

geluInPlace

public void geluInPlace(Tensor a)

tanhInPlace

public void tanhInPlace(Tensor a)

sigmoidInPlace

public void sigmoidInPlace(Tensor a)

multiplyScalar

public Tensor multiplyScalar(Tensor a,
 float scalar)

addScalar

public Tensor addScalar(Tensor a,
 float scalar)

divideScalar

public Tensor divideScalar(Tensor a,
 float scalar)

clamp

public Tensor clamp(Tensor a,
 float min,
 float max)

pow

public Tensor pow(Tensor a,
 float exponent)

sqrt

public Tensor sqrt(Tensor a)

neg

public Tensor neg(Tensor a)

exp

public Tensor exp(Tensor a)

log

public Tensor log(Tensor a)

tanh

public Tensor tanh(Tensor a)

sigmoid

public Tensor sigmoid(Tensor a)

relu

public Tensor relu(Tensor a)

reluBackward

public Tensor reluBackward(Tensor input,
 Tensor gradOutput)

gelu

public Tensor gelu(Tensor a)

geluBackward

public Tensor geluBackward(Tensor input,
 Tensor gradOutput)

softmaxRows

public Tensor softmaxRows(Tensor logits)

softmaxBackward

public Tensor softmaxBackward(Tensor gradOutput,
 Tensor softmaxOut)

softmaxBackward

public Tensor softmaxBackward(Tensor gradient,
 Tensor probabilities,
 float scale)

sumRows

public Tensor sumRows(Tensor a)

sumAlongRows

public Tensor sumAlongRows(Tensor a)

sumAlongCols

public Tensor sumAlongCols(Tensor a)

meanAlongRows

public Tensor meanAlongRows(Tensor a)

varianceAlongRows

public Tensor varianceAlongRows(Tensor a)

maxAlongRows

public Tensor maxAlongRows(Tensor a)

sum

public float sum(Tensor a)

sumAbs

public float sumAbs(Tensor a)

l2Norm

public float l2Norm(List<Tensor> tensors)

transpose

public Tensor transpose(Tensor a)

sumAxes

public Tensor sumAxes(Tensor input,
 int[] axes,
 boolean keep)

meanAxes

public Tensor meanAxes(Tensor input,
 int[] axes,
 boolean keep)

varianceAxes

public Tensor varianceAxes(Tensor input,
 int[] axes,
 boolean keep,
 int correction)

maximumAxes

public Tensor maximumAxes(Tensor input,
 int[] axes,
 boolean keep)

maximumAxesBackward

public Tensor maximumAxesBackward(Tensor input,
 Tensor maximum,
 Tensor gradient,
 int[] axes,
 boolean keep)

softmaxAxis

public Tensor softmaxAxis(Tensor input,
 int axis)

sliceGradient

public Tensor sliceGradient(Tensor gradient,
 int[] shape,
 int axis,
 int start,
 int step)

indexSelect

public Tensor indexSelect(Tensor input,
 int axis,
 int[] indices)

indexAdd

public void indexAdd(Tensor target,
 int axis,
 int[] indices,
 Tensor updates)

indexSelectGradient

public Tensor indexSelectGradient(Tensor gradient,
 int[] shape,
 int axis,
 int[] indices)

isAvailable

public static boolean isAvailable()

supportsNative16Matmul

public static boolean supportsNative16Matmul(DType dtype)

Whether compact matrix products can use native16 operands with FP32 accumulation.

usesDeviceStorage

public boolean usesDeviceStorage()

Description copied from interface: TensorBackend

Whether backend copies reside on a device rather than in ordinary host storage.

Specified by:

usesDeviceStorage in interface TensorBackend

zeros

public Tensor zeros(int rows,
 int columns)

Description copied from interface: TensorBackend

Backend-owned positive-zero storage; ordinary Tensor constructors stay CPU-backed.

Specified by:

zeros in interface TensorBackend

zeros

public Tensor zeros(DType dtype,
 int... shape)

Specified by:

zeros in interface TensorBackend

copy

public Tensor copy(Tensor source)

Description copied from interface: TensorBackend

Independent storage with the same float bits, without recording an autograd operation.

Specified by:

copy in interface TensorBackend

cast

public Tensor cast(Tensor input,
 DType dtype)

Specified by:

cast in interface TensorBackend

copyInto

public void copyInto(Tensor target,
 Tensor source)

Description copied from interface: TensorBackend

Explicitly assigns logical values, casting to the destination's storage dtype.

Specified by:

copyInto in interface TensorBackend

materializeTensor

public void materializeTensor(Tensor t)

Specified by:

materializeTensor in interface TensorBackend

matmul

public Tensor matmul(Tensor a,
 Tensor b)

Specified by:

matmul in interface TensorBackend

sliceRows

public Tensor sliceRows(Tensor input,
 int[] rows)

Specified by:

sliceRows in interface TensorBackend

splitHeads

public Tensor splitHeads(Tensor input,
 int heads)

Specified by:

splitHeads in interface TensorBackend

mergeHeads

public Tensor mergeHeads(Tensor input,
 int heads)

Specified by:

mergeHeads in interface TensorBackend

batchedMatmul

public Tensor batchedMatmul(Tensor left,
 Tensor right,
 int batches,
 boolean transposeLeft,
 boolean transposeRight)

Specified by:

batchedMatmul in interface TensorBackend

causalMask

public Tensor causalMask(Tensor input,
 int sequenceLength)

Specified by:

causalMask in interface TensorBackend

causalSoftmax

public Tensor causalSoftmax(Tensor input,
 int sequenceLength,
 float scale)

Specified by:

causalSoftmax in interface TensorBackend

rotary

public Tensor rotary(Tensor input,
 Tensor cosine,
 Tensor sine,
 int sequenceLength,
 boolean inverse)

Specified by:

rotary in interface TensorBackend

add

public Tensor add(Tensor a,
 Tensor b)

Specified by:

add in interface TensorBackend

subtract

public Tensor subtract(Tensor a,
 Tensor b)

Specified by:

subtract in interface TensorBackend

multiply

public Tensor multiply(Tensor a,
 Tensor b)

Specified by:

multiply in interface TensorBackend

divide

public Tensor divide(Tensor a,
 Tensor b)

Specified by:

divide in interface TensorBackend

addRowVector

public Tensor addRowVector(Tensor a,
 Tensor v)

Specified by:

addRowVector in interface TensorBackend

addBroadcastCols

public Tensor addBroadcastCols(Tensor a,
 Tensor v)

Specified by:

addBroadcastCols in interface TensorBackend

divideBroadcastCols

public Tensor divideBroadcastCols(Tensor a,
 Tensor v)

Specified by:

divideBroadcastCols in interface TensorBackend

subtractBroadcastCols

public Tensor subtractBroadcastCols(Tensor a,
 Tensor v)

Specified by:

subtractBroadcastCols in interface TensorBackend

multiplyBroadcastCols

public Tensor multiplyBroadcastCols(Tensor a,
 Tensor v)

Specified by:

multiplyBroadcastCols in interface TensorBackend

addBroadcastRows

public Tensor addBroadcastRows(Tensor a,
 Tensor v)

Specified by:

addBroadcastRows in interface TensorBackend

multiplyBroadcastRows

public Tensor multiplyBroadcastRows(Tensor a,
 Tensor v)

Specified by:

multiplyBroadcastRows in interface TensorBackend