elements

package
v0.1.0 Latest Latest
Warning

This package is not in the latest version of its module.

 Go to latest Published: Jun 26, 2025 License: Apache-2.0 Imports: 16 Imported by: 1

Documentation

Overview

Package elements provides generic elements, independent of the evaluator, for the interpreter.

Index

Constants

This section is empty.

Variables

This section is empty.

Functions

func AxesFromElement 

func AxesFromElement(el Element) ([]int, error)

AxesFromElement returns a shape from a state element. An error is returned if a concrete shape cannot be returned.

func ConstantFromElement 

func ConstantFromElement(el Element) *values.HostArray

ConstantFromElement returns the host value represented by an element. The function returns (nil, nil) if the element does not host a numerical value.

func ConstantIntFromElement 

func ConstantIntFromElement(el Element) (val int, err error)

ConstantIntFromElement returns a scalar on a host given an element.

func ConstantScalarFromElement 

func ConstantScalarFromElement[T dtype.GoDataType](el Element) (val T, err error)

ConstantScalarFromElement returns a scalar on a host given an element.

func FuncDeclFromElement 

func FuncDeclFromElement(el Element) (*ir.FuncDecl, error)

FuncDeclFromElement extracts a function declaration from an element.

func HostValueFromContext 

func HostValueFromContext(ci *InputValues, el Element) (*values.HostArray, error)

HostValueFromContext returns a host value from the function call.

func ParseCompositeOf 

func ParseCompositeOf[T values.Value](
	f func(ir.Type, []values.Value) (T, error),
) func(ir.Type, []values.Value) (values.Value, error)

ParseCompositeOf returns a function to unflatten a composite value.

func ShapeFromElement 

func ShapeFromElement(node Element) (*shape.Shape, error)

ShapeFromElement returns the shape of a numerical element.

func StringFromElement 

func StringFromElement(el Element) (string, error)

StringFromElement returns the string value stored in a element.

Types

type ArrayOps 

type ArrayOps interface {
	// Graph returns the graph to which new nodes are being added.
	Graph() ops.Graph

	// SubGraph returns a new graph builder.
	SubGraph(name string) (ArrayOps, error)

	// Einsum calls an einstein sum on x and y given the expression in ref.
	Einsum(ref NodeFile[*ir.EinsumExpr], x, y NumericalElement) (NumericalElement, error)

	// BroadcastInDim the data of an array across dimensions.
	BroadcastInDim(expr ExprAt, x NumericalElement, axisLengths []NumericalElement) (NumericalElement, error)

	// Concat concatenates scalars elements into an array with one axis.
	Concat(expr ExprAt, xs []NumericalElement) (NumericalElement, error)

	// Set a slice in an array.
	Set(call NodeFile[*ir.CallExpr], x, updates, index Element) (Element, error)

	// ElementFromArray returns an element from an array GX value.
	ElementFromArray(expr ExprAt, val values.Array) (Node, error)
}

ArrayOps are the operator implementations for arrays.

type ArraySlicer 

type ArraySlicer interface {
	NumericalElement
	SliceArray(ctx FileContext, expr ir.AssignableExpr, index NumericalElement) (NumericalElement, error)
	Type() ir.Type
}

ArraySlicer is a state element with an array that can be sliced.

type CallAt 

type CallAt = NodeFile[*ir.CallExpr]

CallAt is a function call GX expression.

type Copier 

type Copier interface {
	Element
	Copy() Copier
}

Copier is an interface implemented by nodes that need to be copied when passed to a function.

type Element 

type Element interface {
	// Flatten the element, that is:
	// returns itself if the element is atomic,
	// returns its components if the element is a composite.
	Flatten() ([]Element, error)

	// Unflatten creates a GX value from the next handles available in the Unflattener.
	Unflatten(handles *Unflattener) (values.Value, error)

	// Kind returns the kind of the element being stored.
	Kind() ir.Kind
}

Element in the state.

func Flatten 

func Flatten(elts ...Element) ([]Element, error)

Flatten elements.

func Underlying 

func Underlying(val Element) Element

Underlying returns the underlying element.

type ElementWithArrayFromContext 

type ElementWithArrayFromContext interface {
	NumericalElement

	// ArrayFromContext fetches an array from the argument.
	ArrayFromContext(*InputValues) (values.Array, error)
}

ElementWithArrayFromContext is an element able to return a concrete value from the current context. For example, a value passed as an argument to the function.

type ElementWithConstant 

type ElementWithConstant interface {
	NumericalElement

	// NumericalConstant returns the value of a constant represented by a node.
	NumericalConstant() *values.HostArray
}

ElementWithConstant is an element with a concrete value that is already known.

type ExprAt 

type ExprAt = NodeFile[ir.AssignableExpr]

ExprAt is a generic GX expression.

func NewExprAt 

func NewExprAt(file *ir.File, expr ir.AssignableExpr) ExprAt

NewExprAt returns a new expression at a given position.

type FieldAt 

type FieldAt = NodeFile[*ir.Field]

FieldAt is a typed field at a given position.

type FileContext 

type FileContext interface {
	// File in which the operator is being executed.
	File() *ir.File
}

FileContext is context in which an op is being executed.

type Func 

type Func interface {
	Element
	Func() ir.Func
	Recv() *Receiver
	Call(ctx FileContext, call *ir.CallExpr, args []Element) ([]Element, error)
}

Func is an element owning a callable function.

type FuncEvaluator 

type FuncEvaluator interface {
	EvalFunc(f ir.Func, call *ir.CallExpr, args []Element) ([]Element, error)
}

FuncEvaluator is a context able to evaluate a function.

type InputElements 

type InputElements struct {
	// Values are the initial input GX values passed to the function call
	// before they were encapsulated in elements for the interpreter.
	Values InputValues

	// Receiver on which the function call was done.
	// Can be nil.
	Receiver Element

	// Args returns list of arguments passed to the interpreter at call time.
	Args []Element
}

InputElements is the receiver and arguments with which the function was called.

type InputValues 

type InputValues struct {
	// Receiver on which the function call was done.
	// Can be nil.
	Receiver values.Value

	// Args returns list of arguments passed to the interpreter at call time.
	Args []values.Value
}

InputValues are GX values passed to the function call.

type Materialiser 

type Materialiser interface {
	Element

	// Materialise returns the element with all its values from the graph.
	Materialise(ArrayOps) (Node, error)
}

Materialiser is an element that can return an instance of itself composed only of elements from the backend ops.

type NamedType 

type NamedType struct {
	// contains filtered or unexported fields
}

NamedType references a type exported by an imported package.

func NewNamedType 

func NewNamedType(newFunc NewFunc, typ *ir.NamedType, under Copier) *NamedType

NewNamedType returns a new node representing an exported type.

func (*NamedType) Flatten 

func (n *NamedType) Flatten() ([]Element, error)

Flatten returns the named type in a slice of elements.

func (*NamedType) Kind 

func (n *NamedType) Kind() ir.Kind

Kind of the element..

func (*NamedType) NamedType 

func (n *NamedType) NamedType() *ir.NamedType

NamedType returns the type of the element.

func (*NamedType) RecvCopy 

func (n *NamedType) RecvCopy() *NamedType

RecvCopy copies the underlying element and returns the element encapsulated in this named type.

func (*NamedType) Select 

func (n *NamedType) Select(expr SelectAt) (Element, error)

Select returns the field given an index. Returns nil if the receiver type cannot select fields.

func (*NamedType) String 

func (n *NamedType) String() string

String returns a string representation of the node.

func (*NamedType) Unflatten 

func (n *NamedType) Unflatten(handles *Unflattener) (values.Value, error)

Unflatten consumes the next handles to return a GX value.

type NewFunc 

type NewFunc func(ir.Func, *Receiver) Func

NewFunc creates function elements from function IRs.

type Node 

type Node interface {
	Element
	OutNode() *ops.OutputNode
}

Node is an element representing a numerical node in a compute graph.

type NodeAt 

type NodeAt = NodeFile[ir.Node]

NodeAt is a generic GX node.

type NodeFile 

type NodeFile[T ir.Node] struct {
	// contains filtered or unexported fields
}

NodeFile is an expression with the file in which it is declared.

func NewNodeAt 

func NewNodeAt[T ir.Node](file *ir.File, expr T) NodeFile[T]

NewNodeAt returns a new expression at a given position.

func (NodeFile[T]) ExprSrc 

func (ea NodeFile[T]) ExprSrc() ast.Expr

ExprSrc returns the source expression.

func (NodeFile[T]) FSet 

func (ea NodeFile[T]) FSet() *token.FileSet

FSet returns the fileset of the expression.

func (NodeFile[T]) File 

func (ea NodeFile[T]) File() *ir.File

File returns the file in which the expression is declared.

func (NodeFile[T]) Node 

func (ea NodeFile[T]) Node() T

Node returns the expression.

func (NodeFile[T]) NodeFile 

func (ea NodeFile[T]) NodeFile() NodeFile[ir.Node]

NodeFile returns a general node.

func (NodeFile[T]) String 

func (ea NodeFile[T]) String() string

String representation of the node in the source code.

func (NodeFile[T]) ToExprAt 

func (ea NodeFile[T]) ToExprAt() ExprAt

ToExprAt converts a type position into a generic node position.

func (NodeFile[T]) ToNodeAt 

func (ea NodeFile[T]) ToNodeAt() NodeAt

ToNodeAt converts a type position into a generic node position.

func (NodeFile[T]) ToValueAt 

func (ea NodeFile[T]) ToValueAt() ValueAt

ToValueAt converts a type position into a generic node position.

type NumericalElement 

type NumericalElement interface {
	Element

	// UnaryOp applies a unary operator on x.
	UnaryOp(ctx FileContext, expr *ir.UnaryExpr) (NumericalElement, error)

	// BinaryOp applies a binary operator to x and y.
	// Note that the receiver can be either the left or right argument.
	BinaryOp(ctx FileContext, expr *ir.BinaryExpr, x, y NumericalElement) (NumericalElement, error)

	// Cast an element into a given data type.
	Cast(ctx FileContext, expr ir.AssignableExpr, target ir.Type) (NumericalElement, error)

	// Reshape an element.
	Reshape(ctx FileContext, expr ir.AssignableExpr, axisLengths []NumericalElement) (NumericalElement, error)

	// Shape of the value represented by the element.
	Shape() *shape.Shape
}

NumericalElement is a node representing a numerical value.

type Package 

type Package struct {
	// contains filtered or unexported fields
}

Package groups elements exported by a package.

func NewPackage 

func NewPackage(pkg *ir.Package, newFunc NewFunc) *Package

NewPackage returns a package grouping everything that a package exports.

func (*Package) Define 

func (pkg *Package) Define(name string, el Element)

Define an element in the package (used for consts)

func (*Package) Flatten 

func (pkg *Package) Flatten() ([]Element, error)

Flatten returns the package in a slice of

func (*Package) Kind 

func (pkg *Package) Kind() ir.Kind

Kind of the element.

func (*Package) Select 

func (pkg *Package) Select(sel SelectAt) (Element, error)

Select a member of the package.

func (*Package) String 

func (pkg *Package) String() string

String returns a string representation of the node.

func (*Package) Unflatten 

func (pkg *Package) Unflatten(handles *Unflattener) (values.Value, error)

Unflatten creates a GX value from the next handles available in the Unflattener.

type PackageVarSetElement 

type PackageVarSetElement struct {
	// Pck is the package owning the variable.
	Pkg string
	// Index of the variable in the package definition.
	Var string
	// Value of the static variable for the compiler.
	Value Element
}

PackageVarSetElement is an option to set a package variable to an element.

func (PackageVarSetElement) Package 

func (p PackageVarSetElement) Package() string

Package for which the option has been built.

type Receiver 

type Receiver struct {
	Ident   *ast.Ident
	Element *NamedType
}

Receiver of a function.

func NewReceiver 

func NewReceiver(el *NamedType, fn ir.Func) *Receiver

NewReceiver returns a new receiver given a function definition and the element representing the receiver.

type SelectAt 

type SelectAt = NodeFile[*ir.SelectorExpr]

SelectAt is a typed field at a given position.

type Selector 

type Selector interface {
	Select(SelectAt) (Element, error)
}

Selector selects a field given its index.

type Slice 

type Slice struct {
	// contains filtered or unexported fields
}

Slice element storing a slice of elements.

func NewSlice 

func NewSlice(expr ExprAt, elements []Element) *Slice

NewSlice returns a slice from a slice of elements.

func (*Slice) Compare 

func (n *Slice) Compare(x canonical.Comparable) bool

Compare the slice to another canonical value.

func (*Slice) Elements 

func (n *Slice) Elements() []Element

Elements stored in the slice.

func (*Slice) Expr 

func (n *Slice) Expr() (ir.AssignableExpr, error)

Expr returns the IR expression representing the slice.

func (*Slice) Flatten 

func (n *Slice) Flatten() ([]Element, error)

Flatten returns the elements of the slice.

func (*Slice) Kind 

func (*Slice) Kind() ir.Kind

Kind of the element.

func (*Slice) Len 

func (n *Slice) Len() int

Len returns the number of elements in the slice.

func (*Slice) Slice 

func (n *Slice) Slice(ctx FileContext, expr ir.AssignableExpr, index NumericalElement) (Element, error)

Slice of the tuple.

func (*Slice) String 

func (n *Slice) String() string

String returns a string representation of the slice.

func (*Slice) Unflatten 

func (n *Slice) Unflatten(handles *Unflattener) (values.Value, error)

Unflatten consumes the next handles to return a GX value.

type Slicer 

type Slicer interface {
	Slice(ctx FileContext, expr ir.AssignableExpr, index NumericalElement) (Element, error)
}

Slicer is a state element that can be sliced.

type StorageAt 

type StorageAt = NodeFile[ir.Storage]

StorageAt is a generic GX expression.

type String 

type String struct {
	// contains filtered or unexported fields
}

String is a GX string.

func NewString 

func NewString(str *ir.StringLiteral) (*String, error)

NewString returns a state element storing a string GX value.

func (*String) Flatten 

func (n *String) Flatten() ([]Element, error)

Flatten returns the element in a slice of elements.

func (*String) Kind 

func (*String) Kind() ir.Kind

Kind of the element.

func (*String) StringValue 

func (n *String) StringValue() *values.String

StringValue returns the string value as a GX value.

func (*String) Unflatten 

func (n *String) Unflatten(handles *Unflattener) (values.Value, error)

Unflatten consumes the next handles to return a GX value.

type Struct 

type Struct struct {
	// contains filtered or unexported fields
}

Struct is an instance of a structure.

func NewStruct 

func NewStruct(structType *ir.StructType, expr ValueAt, fields map[string]Element) *Struct

NewStruct returns a new node representing a structure instance.

func NewStructFromElements 

func NewStructFromElements(structType *ir.StructType, expr ValueAt, vals []Element) *Struct

NewStructFromElements returns a new node representing a structure instance given a slice of

func (*Struct) Copy 

func (n *Struct) Copy() Copier

Copy the structure to a new node.

func (*Struct) Flatten 

func (n *Struct) Flatten() ([]Element, error)

Flatten returns a flat list of all the elements stored in the structure.

func (*Struct) Kind 

func (*Struct) Kind() ir.Kind

Kind of the element.

func (*Struct) Select 

func (n *Struct) Select(expr SelectAt) (Element, error)

Select returns the value of a field of a structure given its index.

func (*Struct) SetField 

func (n *Struct) SetField(name string, value Element)

SetField sets the field in the structure.

func (*Struct) String 

func (n *Struct) String() string

func (*Struct) StructType 

func (n *Struct) StructType() *ir.StructType

StructType returns the type of the structure.

func (*Struct) Unflatten 

func (n *Struct) Unflatten(handles *Unflattener) (values.Value, error)

Unflatten consumes the next handles to return a GX value.

type Tuple 

type Tuple struct {
	// contains filtered or unexported fields
}

Tuple value grouping multiple values together.

func NewTuple 

func NewTuple(file *ir.File, node ir.Node, values []Element) *Tuple

NewTuple returns a tuple to store the result of a function returning more than one value.

func (*Tuple) Elements 

func (n *Tuple) Elements() []Element

Elements returns the elements stored in the tuple.

func (*Tuple) Flatten 

func (n *Tuple) Flatten() ([]Element, error)

Flatten the tuple and all its elements.

func (*Tuple) Kind 

func (*Tuple) Kind() ir.Kind

Kind of the element.

func (*Tuple) Slice 

func (n *Tuple) Slice(ctx FileContext, expr ir.AssignableExpr, index NumericalElement) (Element, error)

Slice of the tuple.

func (*Tuple) String 

func (n *Tuple) String() string

func (*Tuple) Unflatten 

func (n *Tuple) Unflatten(handles *Unflattener) (values.Value, error)

Unflatten creates a GX value from the next handles available in the Unflattener.

type Unflattener 

type Unflattener struct {
	// contains filtered or unexported fields
}

Unflattener unflattens the output of a graph computation into GX values.

func NewUnflattener 

func NewUnflattener(dev platform.Device, callInputs *InputValues, handles []platform.DeviceHandle) *Unflattener

NewUnflattener returns a new Unflattener given the output of a graph computation.

func (*Unflattener) CallInputs 

func (h *Unflattener) CallInputs() *InputValues

CallInputs returns the inputs with which the function was called.

func (*Unflattener) Device 

func (h *Unflattener) Device() platform.Device

Device returns to which transfers the host value to. TODO(b/388207169): Always transfer the value to device because C++ bindings do not support HostValue.

func (*Unflattener) Next 

func (h *Unflattener) Next() platform.DeviceHandle

Next returns a the next handle and moves the cursor to the succeeding handle.

func (*Unflattener) ParseArray 

func (h *Unflattener) ParseArray(expr ExprAt) (values.Array, error)

ParseArray the next value as an array.

func (*Unflattener) ParseComposite 

func (h *Unflattener) ParseComposite(ncv newCompValue, typ ir.Type, els []Element) (values.Value, error)

ParseComposite unflatten a slice of elements into a single GX value.

func (*Unflattener) String 

func (h *Unflattener) String() string

func (*Unflattener) Unflatten 

func (h *Unflattener) Unflatten(el Element) (values.Value, error)

Unflatten consumes the next available handles and returns a GX value matching the given element.

type ValueAt 

type ValueAt = NodeFile[ir.Value]

ValueAt is a generic GX expression.

func NewValueAt 

func NewValueAt(file *ir.File, expr ir.Value) ValueAt

NewValueAt returns a new expression at a given position.

Source Files

View all Source files

Jump to

Keyboard shortcuts

? : This menu
/ : Search site
f or F : Jump to
y or Y : Canonical URL
go.dev uses cookies from Google to deliver and enhance the quality of its services and to analyze traffic. Learn more.