graph

package

Go to main page

Versions in this module

v1

v1.2.0

Mar 15, 2026

Changes in this version

type ExecutionPlan

+ func (p *ExecutionPlan[T]) EnsureCaptureInputsGPU(start, end int, gpuSlotCache map[int]*tensor.TensorNumeric[T])

+ func (p *ExecutionPlan[T]) PreUploadFrozenWeights() error

type Graph

+ func (g *Graph[T]) AddKVPair(input StatefulInputNode[T], output Node[T])

+ func (g *Graph[T]) ResetStatefulNodes()

type Instruction

+ func FuseRMSNorm[T tensor.Numeric](instructions []Instruction[T], frozenIdx []int, ...) ([]Instruction[T], []int, []*tensor.TensorNumeric[T], [][]int)

+ type Resettable interface

+ Reset func()

+ type StatefulInputNode interface

+ SetStored func(t *tensor.TensorNumeric[T])

v1.1.0

Mar 14, 2026

Changes in this version

+ func ParallelForward[T tensor.Numeric](ctx context.Context, g *Graph[T], inputs ...*tensor.TensorNumeric[T]) (*tensor.TensorNumeric[T], error)

+ type BufferArena struct

+ func NewBufferArena[T tensor.Numeric](shapes [][]int) *BufferArena[T]

+ func (a *BufferArena[T]) Get(idx int) *tensor.TensorNumeric[T]

+ func (a *BufferArena[T]) Len() int

+ func (a *BufferArena[T]) Reset()

+ func (a *BufferArena[T]) Set(idx int, t *tensor.TensorNumeric[T], freeze bool)

+ type BufferLayout struct

+ Offsets []int

+ Sizes []int

+ TotalElements int

+ func ComputeBufferLayout(slotShapes [][]int, frozenIdx []int, inputIdx []int) BufferLayout

+ type CUDAGraphExecutor struct

+ func NewCUDAGraphExecutor[T tensor.Numeric](plan *ExecutionPlan[T], streamPtr unsafe.Pointer, warmups int, ...) *CUDAGraphExecutor[T]

+ func (g *CUDAGraphExecutor[T]) Destroy()

+ func (g *CUDAGraphExecutor[T]) Run(ctx context.Context, inputs ...*tensor.TensorNumeric[T]) (*tensor.TensorNumeric[T], error)

+ type EmbeddedFrozenProvider interface

+ EmbeddedFrozen func() []*tensor.TensorNumeric[T]

+ type ExecutionPlan struct

+ func (p *ExecutionPlan[T]) BufferLayout() *BufferLayout

+ func (p *ExecutionPlan[T]) EnsureSlotsGPU(gpuSlotCache map[int]*tensor.TensorNumeric[T])

+ func (p *ExecutionPlan[T]) FrozenSlots() []FrozenSlot[T]

+ func (p *ExecutionPlan[T]) HasPreallocatedBuffers() bool

+ func (p *ExecutionPlan[T]) InputSlots() []int

+ func (p *ExecutionPlan[T]) InstructionCount() int

+ func (p *ExecutionPlan[T]) InstructionOpName(i int) string

+ func (p *ExecutionPlan[T]) InstructionOutputIdx(i int) int

+ func (p *ExecutionPlan[T]) Instructions() []InstructionMeta

+ func (p *ExecutionPlan[T]) OutputSlot() int

+ func (p *ExecutionPlan[T]) OutputTensor() *tensor.TensorNumeric[T]

+ func (p *ExecutionPlan[T]) PreallocateBuffers()

+ func (p *ExecutionPlan[T]) PrepareSlots(inputs ...*tensor.TensorNumeric[T]) error

+ func (p *ExecutionPlan[T]) Run(ctx context.Context, inputs ...*tensor.TensorNumeric[T]) (*tensor.TensorNumeric[T], error)

+ func (p *ExecutionPlan[T]) RunInstructionRange(ctx context.Context, start, end int) error

+ func (p *ExecutionPlan[T]) RunInstructions(ctx context.Context, inputs ...*tensor.TensorNumeric[T]) (*tensor.TensorNumeric[T], error)

+ func (p *ExecutionPlan[T]) ScratchSlot(idx int) *tensor.TensorNumeric[T]

+ func (p *ExecutionPlan[T]) SetMegakernelFn(...)

+ func (p *ExecutionPlan[T]) SetScratchSlot(idx int, t *tensor.TensorNumeric[T])

+ func (p *ExecutionPlan[T]) SlotShapes() [][]int

+ type FrozenSlot struct

+ Data *tensor.TensorNumeric[T]

+ SlotIdx int

type Graph

+ func FoldConstantTransposes[T tensor.Numeric](g *Graph[T], tr Transposer[T]) (*Graph[T], error)

+ func (g *Graph[T]) ClearMemo()

+ func (g *Graph[T]) Compile(ctx context.Context, inputs ...*tensor.TensorNumeric[T]) (*ExecutionPlan[T], error)

+ func (g *Graph[T]) CompileTraced(ctx context.Context, inputs ...*tensor.TensorNumeric[T]) (*ExecutionPlan[T], error)

+ func (g *Graph[T]) ConstantTensors() []*tensor.TensorNumeric[T]

+ func (g *Graph[T]) Dependencies(n Node[T]) []Node[T]

+ func (g *Graph[T]) EngineProxy() *compute.EngineProxy[T]

+ func (g *Graph[T]) GetAllNodes() []Node[T]

+ func (g *Graph[T]) GetDependencies() map[Node[T]][]Node[T]

+ func (g *Graph[T]) GetNodeMetadata(n Node[T]) map[string]interface{}

+ func (g *Graph[T]) GetTopologicalOrder() ([]Node[T], error)

+ func (g *Graph[T]) Nodes() []Node[T]

+ func (g *Graph[T]) SetEngineProxy(proxy *compute.EngineProxy[T])

+ func (g *Graph[T]) WithParallel(enabled bool)

+ func (g *Graph[T]) WithPool(pool TensorReleaser[T])

+ type Instruction struct

+ ExtraArgs map[string]any

+ Forward func(ctx context.Context, inputs []*tensor.TensorNumeric[T]) (*tensor.TensorNumeric[T], error)

+ InputIdx []int

+ OpName string

+ OutputIdx int

+ type InstructionMeta struct

+ ExtraArgs map[string]any

+ InputIdx []int

+ OpName string

+ OutputIdx int

+ type TensorReleaser interface

+ Release func(t *tensor.TensorNumeric[T])

+ type Transposer interface

+ Transpose func(ctx context.Context, a *tensor.TensorNumeric[T], axes []int, ...) (*tensor.TensorNumeric[T], error)

v0

v0.3.0

Aug 25, 2025

v0.2.0

Aug 22, 2025

Changes in this version

type Graph

+ func (g *Graph[T]) Inputs() []Node[T]

+ func (g *Graph[T]) Output() Node[T]

type Node

+ Attributes func() map[string]interface{}

+ OpType func() string

v0.1.0

Aug 4, 2025

Changes in this version

+ var ErrInvalidInputCount = errors.New("invalid number of input tensors")

+ type Builder struct

+ func NewBuilder[T tensor.Numeric](engine compute.Engine[T]) *Builder[T]

+ func (b *Builder[T]) AddNode(node Node[T], inputs ...Node[T]) Node[T]

+ func (b *Builder[T]) Build(outputNode Node[T]) (func(inputs ...*tensor.Tensor[T]) (*tensor.Tensor[T], error), ...)

+ func (b *Builder[T]) Input(shape []int) Node[T]

+ func (b *Builder[T]) Parameters() []*Parameter[T]

+ type Graph struct

+ func (g *Graph[T]) Backward(initialGradient *tensor.Tensor[T]) error

+ func (g *Graph[T]) Forward(inputs ...*tensor.Tensor[T]) (*tensor.Tensor[T], error)

+ func (g *Graph[T]) Parameters() []*Parameter[T]

+ type NoParameters struct

+ func (np *NoParameters[T]) Parameters() []*Parameter[T]

+ type Node interface

+ Backward func(outputGradient *tensor.Tensor[T]) ([]*tensor.Tensor[T], error)

+ Forward func(inputs ...*tensor.Tensor[T]) (*tensor.Tensor[T], error)

+ OutputShape func() []int

+ Parameters func() []*Parameter[T]

+ type Parameter struct

+ Gradient *tensor.Tensor[T]

+ Name string

+ Value *tensor.Tensor[T]

+ func NewParameter[T tensor.Numeric](name string, value *tensor.Tensor[T], ...) (*Parameter[T], error)

+ func (p *Parameter[T]) AddGradient(grad *tensor.Tensor[T]) error

+ func (p *Parameter[T]) ClearGradient()

?	: This menu
/	: Search site
f or F	: Jump to
y or Y	: Canonical URL