Package io.github.kirstenali.deepj.autograd

Class LayoutOps

java.lang.Object

io.github.kirstenali.deepj.autograd.LayoutOps

public final class LayoutOps
extends Object

Method Summary

Modifier and Type	Method	Description
static Tensor	batchedMatmul(Tensor a, Tensor b, int batches, boolean ta, boolean tb)
static Tensor	gather(Tensor input, int[] indices)
static Tensor	indexSelect(Tensor input, int axis, int[] indices)
static Tensor	mergeHeads(Tensor input, int heads)
static Tensor	rotary(Tensor input, Tensor cosine, Tensor sine, int length, boolean inverse)
static Tensor	splitHeads(Tensor input, int heads)

Methods inherited from class java.lang.Object

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

Method Details

indexSelect

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

gather

public static Tensor gather(Tensor input,
 int[] indices)

splitHeads

public static Tensor splitHeads(Tensor input,
 int heads)

mergeHeads

public static Tensor mergeHeads(Tensor input,
 int heads)

batchedMatmul

public static Tensor batchedMatmul(Tensor a,
 Tensor b,
 int batches,
 boolean ta,
 boolean tb)

rotary

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