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 a Tensor's unit element.
Shape and DType are used both by the concrete tensor values (see tensor package) and when working on the computation graph (see graph package).
Go float16 support (commonly used by Nvidia GPUs) uses github.com/x448/float16 implementation, and bfloat16 uses a simple implementation in github.com/gomlx/gopjrt/dtypes/bfloat16.
## Glossary
- Rank: number of axes (dimensions) of a Tensor.
- Axis: is the index of a dimension on a multidimensional 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-dimension Tensor in one of its axes. See the example below.
- DType: the data type of the unit element in a tensor. Enumeration defined in github.com/gomlx/gopjrt/dtypes
- 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 graph nodes -- 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_ functionality. Examples:
AssertRank and AssertDims check 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 look like `logits.AssertDims(batchSize, -1)`.
Index ¶
- Constants
- func Assert(shaped HasShape, dtype dtypes.DType, dimensions ...int)
- func AssertDims(shaped HasShape, dimensions ...int)
- func AssertRank(shaped HasShape, rank int)
- func AssertScalar(shaped HasShape)
- func CastAsDType(value any, dtype dtypes.DType) any
- func CheckDims(shaped HasShape, dimensions ...int) error
- func CheckRank(shaped HasShape, rank int) error
- func CheckScalar(shaped HasShape) error
- func ConvertTo[T dtypes.NumberNotComplex](value any) T
- func UnsafeSliceForDType(dtype dtypes.DType, unsafePtr unsafe.Pointer, len int) (any, error)
- type HasShape
- type Shape
- func ConcatenateDimensions(s1, s2 Shape) (shape Shape)
- func FromAnyValue(v any) (shape Shape, err error)
- func GobDeserialize(decoder *gob.Decoder) (s Shape, err error)
- func Invalid() Shape
- func Make(dtype dtypes.DType, dimensions ...int) Shape
- func MakeTuple(elements []Shape) Shape
- func Scalar[T dtypes.Number]() Shape
- func (s Shape) Assert(dtype dtypes.DType, dimensions ...int)
- func (s Shape) AssertDims(dimensions ...int)
- func (s Shape) AssertRank(rank int)
- func (s Shape) AssertScalar()
- func (s Shape) Check(dtype dtypes.DType, dimensions ...int) error
- func (s Shape) CheckDims(dimensions ...int) error
- func (s Shape) CheckRank(rank int) error
- func (s Shape) CheckScalar() error
- func (s Shape) Clone() (s2 Shape)
- func (s Shape) Dim(axis int) int
- func (s Shape) Equal(s2 Shape) bool
- func (s Shape) EqualDimensions(s2 Shape) bool
- func (s Shape) GobSerialize(encoder *gob.Encoder) (err error)
- func (s Shape) IsScalar() bool
- func (s Shape) IsTuple() bool
- func (s Shape) IsZeroSize() bool
- func (s Shape) Memory() uintptr
- func (s Shape) Ok() bool
- func (s Shape) Rank() int
- func (s Shape) Shape() Shape
- func (s Shape) Size() (size int)
- func (s Shape) String() string
- func (s Shape) ToStableHLO() string
- func (s Shape) TupleSize() int
- func (s Shape) WriteStableHLO(writer io.Writer) error
Constants ¶
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 Assert ¶
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 AssertDims ¶
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 ¶
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 CastAsDType ¶
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 ¶
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 ¶
CheckRank checks that the shape has the given rank.
It returns an error if the rank is different.
func CheckScalar ¶
CheckScalar checks that the shape is a scalar.
It returns an error if shape is not a scalar.
func ConvertTo ¶
func ConvertTo[T dtypes.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.
Types ¶
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 Shape ¶
type Shape struct {
DType dtypes.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 ¶
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 FromAnyValue ¶
FromAnyValue attempts to convert a Go "any" value to its expected shape. Accepted values are plain-old-data (POD) types (ints, floats, complex), slices (or multiple level of slices) of POD.
It returns the expected shape.
Example:
shape := shapes.FromAnyValue([][]float64{{0, 0}}) // Returns shape (Float64)[1 2]
func GobDeserialize ¶
GobDeserialize a Shape. Returns new Shape or an error.
func Make ¶
Make returns a Shape structure filled with the values given. See MakeTuple for tuple shapes.
func (Shape) Assert ¶
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 ¶
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 ¶
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 ¶
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 ¶
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 ¶
CheckRank checks that the shape has the given rank.
It returns an error if the rank is different.
func (Shape) CheckScalar ¶
CheckScalar checks that the shape is a scalar.
It returns an error if shape is not a scalar.
func (Shape) Dim ¶
Dim returns the dimension of the given axis. axis can take negative numbers, in which case it counts as starting from the end -- so axis=-1 refers to the last axis. Like with a slice indexing, it panics for an out-of-bound axis.
func (Shape) EqualDimensions ¶
EqualDimensions compares two shapes for equality of dimensions. Dtypes can be different.
func (Shape) GobSerialize ¶
GobSerialize shape in binary format.
func (Shape) IsScalar ¶
IsScalar returns whether the shape represents a scalar, that is there are no dimensions (rank==0).
func (Shape) IsZeroSize ¶
IsZeroSize returns whether any of the dimensions is zero, in which case it's an empty shape, with no data attached to it.
Notice scalars are not zero in size -- they have size one, but rank zero.
func (Shape) Memory ¶
Memory returns the memory used to store an array of the given shape, the same as the size in bytes. Careful, so far all types in Go and on device seem to use the same sizes, but future type this is not guaranteed.
func (Shape) Ok ¶
Ok returns whether this is a valid Shape. A "zero" shape, that is just instantiating it with Shape{} will be invalid.
func (Shape) Size ¶
Size returns the number of elements (not bytes) for this shape. It's the product of all dimensions.
For the number of bytes used to store this shape, see Shape.Memory.
func (Shape) ToStableHLO ¶
ToStableHLO returns the ToStableHLO representation of the shape's type.
Source Files