shapes

package
v0.9.1 Latest Latest
Warning

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

 Go to latest Published: Apr 20, 2024 License: Apache-2.0 Imports: 11 Imported by: 0

Documentation

Overview

Package shapes defines Shape and DType and associated tools.

Shape represents the shape (rank, dimensions and DType) of either a Tensor or the expected shape of a node in a computation Graph. DType indicates the type of the unit element of a Tensor (or its representation as a node in a computation Graph).

Shape and DType are used both by the concrete tensor values (see tensor package) and when working on the computation graph (see graph package).

## Glossary

  • Rank: number of axes (dimensions) of a Tensor.
  • Axis: is the index of a dimension on a multi-dimensional Tensor. Sometimes used interchangeably with Dimension, but here we try to refer to a dimension index as "axis" (plural axes), and its size as its dimension.
  • Dimension: the size of a multi-dimensions Tensor in one of its axes. See example below:
  • DType: the data type of the unit element in a tensor.
  • Scalar: is a shape where there are no axes (or dimensions), only a single value of the associated DType.

Example: The multi-dimensional array `[][]int32{{0, 1, 2}, {3, 4, 5}}` if converted to a Tensor would have shape `(int32)[2 3]`. We say it has rank 2 (so 2 axes), axis 0 has dimension 2, and axis 1 has dimension 3. This shape could be created with `shapes.Make(int32, 2, 3)`.

## Asserts

When coding ML models, one delicate part is keeping tabs on the shape of the nodes of the graphs -- unfortunately there is no compile-time checking of values, so validation only happens in runtime. To facilitate, and also to serve as code documentation, this package provides two variations of _assert_ funtionality. Examples:

`AssertRank` and `AssertDims` checks that the rank and dimensions of the given 

object (that has a `Shape` method) match, otherwise it panics. The `-1` means
the dimension is unchecked (it can be anything).

``` 

func modelGraph(ctx *context.Context, spec any, inputs []*Node) ([]*Node) {
   _ = spec  // Not needed here, we know the dataset.
   shapes.AssertRank(inputs, 2)
   batchSize := inputs.Shape().Dimensions[0]
   logits := layers.Dense(ctx, inputs[0], /* useBias= */ true, /* outputDim= */ 1)
   shapes.AssertDims(logits, batchSize, -1)
   return []*Node{logits}
}

```

If you don't want to panic, but instead return an error through the `graph.Graph`, you can use the `Node.AssertDims()` method. So it would loook like `logits.AssertDims(batchSize, -1)`.

Index

Constants

View Source 
const (
	U8   = UInt8
	U32  = UInt32
	U64  = UInt64
	I32  = Int32
	I64  = Int64
	F16  = Float16
	F32  = Float32
	F64  = Float64
	C64  = Complex64
	C128 = Complex128
)
View Source 
const PRED = Bool

PRED type is an alias to Bool, used in `tensorflow/compiler/xla/xla_data.proto`.

View Source 
const UncheckedAxis = int(-1)

UncheckedAxis can be used in CheckDims or AssertDims functions for an axis whose dimension doesn't matter.

Variables

This section is empty.

Functions

func AssertDims 

func AssertDims(shaped HasShape, dimensions ...int)

AssertDims checks that the shape has the given dimensions and rank. A value of -1 in dimensions means it can take any value and is not checked.

It panics if it doesn't match.

See usage example in package shapes documentation.

func AssertRank 

func AssertRank(shaped HasShape, rank int)

AssertRank checks that the shape has the given rank.

It panics if it doesn't match.

See usage example in package shapes documentation.

func AssertScalar 

func AssertScalar(shaped HasShape)

AssertScalar checks that the shape is a scalar.

It panics if it doesn't match.

See usage example in package shapes documentation.

func AssertShape added in v0.9.0 

func AssertShape(dtype DType, dimensions ...int)

func CastAsDType 

func CastAsDType(value any, dtype DType) any

CastAsDType casts a numeric value to the corresponding for the DType. If the value is a slice it will convert to a newly allocated slice of the given DType.

It doesn't work for complex numbers.

func CheckDims 

func CheckDims(shaped HasShape, dimensions ...int) error

CheckDims checks that the shape has the given dimensions and rank. A value of -1 in dimensions means it can take any value and is not checked.

It returns an error if the rank is different or any of the dimensions.

func CheckRank 

func CheckRank(shaped HasShape, rank int) error

CheckRank checks that the shape has the given rank.

It returns an error if the rank is different.

func CheckScalar 

func CheckScalar(shaped HasShape) error

CheckScalar checks that the shape is a scalar.

It returns an error if shape is not a scalar.

func ConvertTo added in v0.2.1 

func ConvertTo[T NumberNotComplex](value any) T

ConvertTo converts any scalar (typically returned by `tensor.Local.Value()`) of the supported dtypes to `T`. Returns 0 if value is not a scalar or not a supported number (e.g: bool). It doesn't work for if T (the output type) is a complex number. If value is a complex number, it converts by taking the real part of the number and discarding the imaginary part.

func DTypeStrings added in v0.4.0 

func DTypeStrings() []string

DTypeStrings returns a slice of all String values of the enum

func LowestValueForDType 

func LowestValueForDType(dtype DType) any

LowestValueForDType converted to the corresponding Go type. For float values it will return negative infinite. There is no lowest value for complex numbers, since they are not ordered.

func SmallestNonZeroValueForDType added in v0.4.1 

func SmallestNonZeroValueForDType(dtype DType) any

SmallestNonZeroValueForDType is the smallest non-zero value dtypes. Only useful for float types. The return value is converted to the corresponding Go type. There is no smallest non-zero value for complex numbers, since they are not ordered.

func TypeForDType 

func TypeForDType(dtype DType) reflect.Type

TypeForDType returns the Go `reflect.Type` corresponding to the tensor DType.

func UnsafeSliceForDType added in v0.2.1 

func UnsafeSliceForDType(dtype DType, unsafePtr unsafe.Pointer, len int) any

UnsafeSliceForDType creates a slice of the corresponding dtype and casts it to any. It uses unsafe.Slice. Set `len` to the number of `DType` elements (not the number of bytes).

Types

type DType 

type DType int32

DType indicates the type of the unit element of a Tensor (or its representation in a computation graph). It enumerates the known data types. So far only Bool, Uint8 (U8), Int32 (I32), Int64 (I64), Uint64 (U64), Float32 (F32) and Float64 (F64) are supported.

The values of DType must match "tensorflow/compiler/xla/xla_data.pb.h", hence it needs to be an int32.

See example in package shapes documentation. 

const (
	InvalidDType DType = iota
	Bool               // Bool, but also known as PRED in `xla_data.proto`.
	Int8               // S8
	Int16              // S16
	Int32              // S32
	Int64              // S64
	UInt8              // U8
	UInt16             // U16
	UInt32             // U32
	UInt64             // U64
	Float16            // F16
	Float32            // F32
	Float64            // F64

	BFloat16   DType = 16 // BF16
	Complex64  DType = 15 // C64
	Complex128 DType = 18 // C128

	Tuple      DType = 13
	OpaqueType DType = 14
	Token      DType = 17
)

DType constants must match `tensorflow/compiler/xla/xla_data.proto`.

func DTypeForType 

func DTypeForType(t reflect.Type) DType

func DTypeGeneric 

func DTypeGeneric[T Supported]() DType

func DTypeString added in v0.4.0 

func DTypeString(s string) (DType, error)

DTypeString retrieves an enum value from the enum constants string name. Throws an error if the param is not part of the enum.

func DTypeValues added in v0.4.0 

func DTypeValues() []DType

DTypeValues returns all values of the enum

func (DType) GoStr added in v0.7.2 

func (dtype DType) GoStr() string

GoStr converts dtype to the corresponding Go type and convert that to string. Notice the names are different from the Dtype (so `Int64` dtype is simply `int` in Go).

func (DType) IsADType added in v0.4.0 

func (i DType) IsADType() bool

IsADType returns "true" if the value is listed in the enum definition. "false" otherwise

func (DType) IsComplex added in v0.6.0 

func (dtype DType) IsComplex() bool

IsComplex returns whether dtype is a supported complex number type.

func (DType) IsFloat 

func (dtype DType) IsFloat() bool

IsFloat returns whether dtype is a supported float -- float types not yet supported will return false. It returns false for complex numbers.

func (DType) IsInt 

func (dtype DType) IsInt() bool

IsInt returns whether dtype is a supported integer type -- float types not yet supported will return false.

func (DType) IsSupported 

func (dtype DType) IsSupported() bool

func (DType) MarshalJSON added in v0.4.0 

func (i DType) MarshalJSON() ([]byte, error)

MarshalJSON implements the json.Marshaler interface for DType

func (DType) MarshalText added in v0.4.0 

func (i DType) MarshalText() ([]byte, error)

MarshalText implements the encoding.TextMarshaler interface for DType

func (DType) MarshalYAML added in v0.4.0 

func (i DType) MarshalYAML() (interface{}, error)

MarshalYAML implements a YAML Marshaler for DType

func (DType) Memory added in v0.2.0 

func (dtype DType) Memory() int64

Memory returns the number of bytes in Go used to store the given DType.

func (DType) RealDType added in v0.6.0 

func (dtype DType) RealDType() DType

RealDType returns the real component of complex dtypes. For float dtypes, it returns itself.

It returns InvalidDType for other non-(complex or float) dtypes.

func (DType) String 

func (i DType) String() string

func (DType) Type added in v0.2.0 

func (dtype DType) Type() reflect.Type

Type returns the Go `reflect.Type` corresponding to the tensor DType.

func (*DType) UnmarshalJSON added in v0.4.0 

func (i *DType) UnmarshalJSON(data []byte) error

UnmarshalJSON implements the json.Unmarshaler interface for DType

func (*DType) UnmarshalText added in v0.4.0 

func (i *DType) UnmarshalText(text []byte) error

UnmarshalText implements the encoding.TextUnmarshaler interface for DType

func (*DType) UnmarshalYAML added in v0.4.0 

func (i *DType) UnmarshalYAML(unmarshal func(interface{}) error) error

UnmarshalYAML implements a YAML Unmarshaler for DType

func (DType) Values added in v0.4.0 

func (DType) Values() []string

type GoFloat 

type GoFloat interface {
	float32 | float64
}

GoFloat represent a continuous Go numeric type, supported by GoMLX.

type HasShape 

type HasShape interface {
	Shape() Shape
}

HasShape is an interface for objects that have an associated Shape. `tensor.Tensor` (concrete tensor) and `graph.Node` (tensor representations in a computation graph), `context.Variable` and Shape itself implement the interface.

type MultiDimensionSlice 

type MultiDimensionSlice interface {
	bool | float32 | float64 | int | int32 | int64 | uint8 | uint32 | uint64 | complex64 | complex128 |
		[]bool | []float32 | []float64 | []int | []int32 | []int64 | []uint8 | []uint32 | []uint64 | []complex64 | []complex128 |
		[][]bool | [][]float32 | [][]float64 | [][]int | [][]int32 | [][]int64 | [][]uint8 | [][]uint32 | [][]uint64 | [][]complex64 | [][]complex128 |
		[][][]bool | [][][]float32 | [][][]float64 | [][][]int | [][][]int32 | [][][]int64 | [][][]uint8 | [][][]uint32 | [][][]uint64 | [][][]complex64 | [][][]complex128 |
		[][][][]bool | [][][][]float32 | [][][][]float64 | [][][][]int | [][][][]int32 | [][][][]int64 | [][][][]uint8 | [][][][]uint32 | [][][][]uint64 | [][][][]complex64 | [][][][]complex128 |
		[][][][][]bool | [][][][][]float32 | [][][][][]float64 | [][][][][]int | [][][][][]int32 | [][][][][]int64 | [][][][][]uint8 | [][][][][]uint32 | [][][][][]uint64 | [][][][][]complex64 | [][][][][]complex128 |
		[][][][][][]bool | [][][][][][]float32 | [][][][][][]float64 | [][][][][][]int | [][][][][][]int32 | [][][][][][]int64 | [][][][][][]uint8 | [][][][][][]uint32 | [][][][][][]uint64 | [][][][][][]complex64 | [][][][][][]complex128
}

MultiDimensionSlice lists the Go types a Tensor can be converted to/from. There are no recursions in generics' constraint definitions, so we enumerate up to 7 levels of slices. Feel free to add more if needed, the implementation will work with any arbitrary number.

type Number 

type Number interface {
	float32 | float64 | int | int32 | int64 | uint8 | uint32 | uint64 | complex64 | complex128
}

Number represents the Go numeric types that are supported by graph package. Used as a Generics constraint. Notice that "int" becomes int64 in the implementation. Since it needs a 1:1 mapping, it doesn't support the native (Go) int64 type. It includes complex numbers.

type NumberNotComplex added in v0.6.0 

type NumberNotComplex interface {
	float32 | float64 | int | int32 | int64 | uint8 | uint32 | uint64
}

NumberNotComplex represents the Go numeric types that are supported by graph package except the complex numbers. Used as a Generics constraint. See Number for details.

type Shape 

type Shape struct {
	DType       DType
	Dimensions  []int
	TupleShapes []Shape // Shapes of the tuple, if this is a tuple.
}

Shape represents the shape of either a Tensor or the expected shape of the value from a computation node.

Use Make to create a new shape. See example in package shapes documentation.

func ConcatenateDimensions 

func ConcatenateDimensions(s1, s2 Shape) (shape Shape)

ConcatenateDimensions of two shapes. The resulting rank is the sum of both ranks. They must have the same dtype. If any of them is a scalar, the resulting shape will be a copy of the other. It doesn't work for Tuples.

func GobDeserialize added in v0.2.0 

func GobDeserialize(decoder *gob.Decoder) (s Shape, err error)

GobDeserialize a Shape. Returns new Shape or an error.

func Make 

func Make(dtype DType, dimensions ...int) Shape

Make returns a Shape structure filled with the values given.

func MakeTuple 

func MakeTuple(elements []Shape) Shape

MakeTuple returns a shape representing a tuple of elements with the given shapes.

func Scalar 

func Scalar[T Number]() Shape

Scalar returns a scalar Shape for the given type.

func (Shape) Assert added in v0.9.0 

func (s Shape) Assert(dtype DType, dimensions ...int)

Assert checks that the shape has the given dtype, dimensions and rank. A value of -1 in dimensions means it can take any value and is not checked.

It panics if it doesn't match.

func (Shape) AssertDims 

func (s Shape) AssertDims(dimensions ...int)

AssertDims checks that the shape has the given dimensions and rank. A value of -1 in dimensions means it can take any value and is not checked.

It panics if it doesn't match.

See usage example in package shapes documentation.

func (Shape) AssertRank 

func (s Shape) AssertRank(rank int)

AssertRank checks that the shape has the given rank.

It panics if it doesn't match.

See usage example in package shapes documentation.

func (Shape) AssertScalar 

func (s Shape) AssertScalar()

AssertScalar checks that the shape is a scalar.

It panics if it doesn't match.

See usage example in package shapes documentation.

func (Shape) Check added in v0.9.0 

func (s Shape) Check(dtype DType, dimensions ...int) error

Check that the shape has the given dtype, dimensions and rank. A value of -1 in dimensions means it can take any value and is not checked.

It returns an error if the dtype or rank is different or if any of the dimensions don't match.

func (Shape) CheckDims 

func (s Shape) CheckDims(dimensions ...int) error

CheckDims checks that the shape has the given dimensions and rank. A value of -1 in dimensions means it can take any value and is not checked.

It returns an error if the rank is different or if any of the dimensions don't match.

func (Shape) CheckRank 

func (s Shape) CheckRank(rank int) error

CheckRank checks that the shape has the given rank.

It returns an error if the rank is different.

func (Shape) CheckScalar 

func (s Shape) CheckScalar() error

CheckScalar checks that the shape is a scalar.

It returns an error if shape is not a scalar.

func (Shape) Copy 

func (s Shape) Copy() (s2 Shape)

Copy makes a deep copy of the shapes.

func (Shape) Eq 

func (s Shape) Eq(s2 Shape) bool

Eq compares two shapes for equality: dtype and dimensions are compared.

func (Shape) EqDimensions added in v0.9.0 

func (s Shape) EqDimensions(s2 Shape) bool

EqDimensions compares two shapes for equality of dimensions. Dtypes can be different.

func (Shape) GobSerialize added in v0.2.0 

func (s Shape) GobSerialize(encoder *gob.Encoder) (err error)

GobSerialize shape in binary format.

func (Shape) IsScalar 

func (s Shape) IsScalar() bool

IsScalar returns whether the shape represents a scalar, that is there are no dimensions (rank==0).

func (Shape) IsTuple 

func (s Shape) IsTuple() bool

IsTuple returns whether the shape represents a tuple.

func (Shape) Memory added in v0.2.0 

func (s Shape) Memory() int64

Memory returns the number of bytes for that would be used in Go to store the given data -- the actual memory may depend on the device implementation in some cases (e.g. bool).

func (Shape) Ok 

func (s Shape) Ok() bool

Ok returns whether this is a valid Shape. A "zero" shape, that is just instantiating it with Shape{} will be invalid.

func (Shape) Rank 

func (s Shape) Rank() int

Rank of the shape, that is, the number of dimensions.

func (Shape) Shape 

func (s Shape) Shape() Shape

Shape returns a shallow copy of itself. It implements the HasShape interface.

func (Shape) Size 

func (s Shape) Size() (size int)

Size returns the number of elements of DType are needed for this shape. It's the product of all dimensions.

func (Shape) String 

func (s Shape) String() string

String implements stringer, pretty-prints the shape.

func (Shape) TupleSize 

func (s Shape) TupleSize() int

TupleSize returns the number of elements in the tuple, if it is a tuple.

type Supported 

type Supported interface {
	bool | float32 | float64 | int | int32 | int64 | uint8 | uint32 | uint64 | complex64 | complex128
}

Supported lists the Go types that are supported by the graph package. Used as a Generics constraint. See also Number.

Notice Go's `int` type is not portable, since it may translate to dtypes Int32 or Int64 depending on the platform.

Generated by `cmd/constraints_generator`.

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.