tensors

package
v0.20.1 Latest Latest
Warning

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

Go to latest
Published: Jun 12, 2025 License: Apache-2.0 Imports: 18 Imported by: 13

Documentation

Overview

Package tensors implements a `Tensor`, a representation of a multi-dimensional array.

Tensors are multidimensional arrays (from scalar with 0 dimensions, to arbitrarily large dimensions), defined by their shape (a data type and its axes' dimensions) and their actual content. As a special case, a Tensor can also be a tuple of multiple tensors.

The main use of tensors are to be used as input and output of GoMLX computation graphs.

There are various ways to construct a Tensor from local data:

  • FromShape(shape shapes.Shape): creates a tensor with the given shape, and zero values.

  • FromScalarAndDimensions[T shapes.Supported](value T, dimensions ...int): creates a Tensor with the given dimensions, filled with the scalar value given. `T` must be one of the supported types.

  • FromFlatDataAndDimensions[T shapes.Supported](data []T, dimensions ...int): creates a Tensor with the given dimensions, and set the flattened values with the given data. `T` must be one of the supported types. Example:

    t := FromFlatDataAndDimensions([]int8{1, 2, 3, 4}, 2, 2}) // Tensor with [[1,2], [3,4]]

  • FromValue[S MultiDimensionSlice](value S): Generic conversion, works with the scalar supported `DType`s as well as with any arbitrary multidimensional slice of them. Slices of rank > 1 must be regular, that is all the sub-slices must have the same shape. Example:

    t := FromValue([][]float{{1,2}, {3, 5}, {7, 11}})`

  • FromAnyValue(value any): same as FromValue but non-generic, it takes an anonymous type `any`. The exception is if `value` is already a tensor, then it is a no-op, and it returns the tensor itself.

Behind the scenes (as much as possible Tensor tries to hide all the details), Tensor is a container that keeps in sync different materialization's of value:

  • `local`: a copy of the values stored in CPU, as a Go flat array of the underlying dtype.
  • `onDevices`: a copy of the values stored in the accelerator device(s) (CPU, GPU, TPU, etc.), a wrapper for whatever the backend uses as buffer managed by the lower levels (see github.com/gomlx/gopjrt for the XLA backend). There can be multiple `Device` backing of a tensor if there are multiple devices (like a multi-GPU set up).
  • And "on-device" Tensor can also be "shared" if the backend allows it, in which case the local and "on-device" share the same memory allocation.

The Tensor container is lazy in nature: it won't transfer data from local storage to "on device" until needed. And if/when it can, it will make it "shared" (generally, when running on CPUs). If not "shared", when one (local or on-device) is updated, the others are immediately invalidated.

Transferring tensors to/from local/device areas has a cost and should be avoided. For example, while training weights of an ML model, one generally does not need to transfer those weights to local -- just at the end of training to save the model weights. But the Tensor will keep the (local/device) copies cached, so they can be used multiple times, and transfer only occurs once.

Tensors used across multiple backends:

There is not an easy interface to share tensors across backends (instances) yet, even if they are the same type of backends (if you created two instances of `xla:cpu` backend, for instance).

The recommendation is to keep tensors used for each backend in separate variables and copy when needed: You can use `Tensor.LocalClone()` to copy a tensor from one backend to a new local tensor. Or you can use `Tensor.OnDeviceClone()` to copy a tensor from one backend directly into another backend.

Alternatively, after using a tensor as input to a Backend computation, or a tensor returned from the Backend, call Tensor.ToLocal(): it will remove any links to the backend (by copying all the data locally) and it can then be used by other backends.

Index

Constants

This section is empty.

Variables

View Source
var MaxSizeForString = 500

MaxSizeForString is the largest Local tensor that is actually returned by String() is requested.

Functions

func AssignFlatData added in v0.16.0

func AssignFlatData[T dtypes.Supported](toTensor *Tensor, fromFlat []T)

AssignFlatData will copy over the values in fromFlat to the storage used by toTensor. If the dtypes are not compatible or if the size is wrong, it will panic.

func ConstFlatData

func ConstFlatData[T dtypes.Supported](t *Tensor, accessFn func(flat []T))

ConstFlatData calls accessFn with the flattened data as a slice of the Go type corresponding to the DType type. Even scalar values have a flattened data representation of one element. It locks the Tensor until accessFn returns.

It is the "generics" version of Tensor.ConstFlatData(),

This provides accessFn with the actual Tensor data (not a copy), and it's owned by the Tensor, but it should not be changed -- the contents of the corresponding "on device" tensors could go out-of-sync. See Tensor.MutableFlatData to access a mutable version of the flat data.

See Tensor.Size for the number of elements, and Tensor.LayoutStrides to calculate the offset of individual positions, given the indices at each axis.

It panics if the tensor is in an invalid state (if it was finalized), or if it is a tuple.

func CopyFlatData

func CopyFlatData[T dtypes.Supported](t *Tensor) []T

CopyFlatData returns a copy of the flat data of the Tensor.

It triggers a synchronous transfer from device to local, if the tensor is only on device.

It will panic if the given generic type doesn't match the DType of the tensor.

func MutableFlatData

func MutableFlatData[T dtypes.Supported](t *Tensor, accessFn func(flat []T))

MutableFlatData calls accessFn with a flat slice pointing to the Tensor data. The type of the slice is corresponds to the DType of the tensor. The contents of the slice itself can be changed until accessFn returns. During this time the Tensor is locked.

It is the "generics" version of Tensor.MutableFlatData(), see its description for more details.

This returns the actual Tensor data (not a copy), and the data owned by the Tensor, and should only be changed inside accessFn.

It panics if the tensor is in an invalid state (if it was finalized), or if it is a tuple.

func ToScalar

func ToScalar[T dtypes.Supported](t *Tensor) T

ToScalar returns the scalar value of the Tensor.

It triggers a synchronous transfer from device to local, if the tensor is only on device.

It will panic if the given generic type doesn't match the DType of the tensor.

Types

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.

Generated by `github.com/gomlx/gomlx/cmd/constraints_generator`.

type Tensor

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

Tensor represents a multidimensional array (from scalar with 0 dimensions, to arbitrarily large dimensions), defined by their shape, a data type (dtypes.DType) and its axes' dimensions, and their actual content stored as a flat (1D) array of values.

The main use of tensors is to be used as GoMLX computation graphs inputs and outputs.

It is a container for "local" (host CPU) and "on-device" backing of the tensor -- they can be the same ("shared") in some cases. It is always stored as a flat slice of the underlying DType.

Tensor manages caching of Local and Device copies. There is a transferring cost that one needs to be aware of when using it for large data -- LLM models can have hundreds of GB in size. There is a cache system to prevent duplicate transfers, but it requires some care from the user (see ConstFlatData and MutableFlatData).

Tensors used across multiple backends:

There is not an easy interface to share tensors across backends (instances) yet, even if they are the same type of backends (if you created two instances of `xla:cpu` backend, for instance).

The recommendation is to keep tensors used for each backend in separate variables and copy when needed: You can use `Tensor.LocalClone()` to copy a tensor from one backend to a new local tensor. Or you can use `Tensor.OnDeviceClone()` to copy a tensor from one backend directly into another backend.

Alternatively, after using a tensor as input to a Backend computation, or a tensor returned from the Backend, call Tensor.ToLocal(): it will remove any links to the backend (by copying all the data locally) and it can then be used by other backends.

More details in the `tensor` package documentation.

func FromAnyValue

func FromAnyValue(value any) (t *Tensor)

FromAnyValue is a non-generic version of FromValue that returns a *tensors.Tensor (not specified if local or on device). The input is expected to be either a scalar or a slice of slices with homogeneous dimensions. If the input is a tensor already (Local or Device), it is simply returned. If value is anything but a Device tensor, it will return a Local tensor.

It panics with an error if `value` type is unsupported or the shape is not regular.

func FromBuffer

func FromBuffer(backend backends.Backend, buffer backends.Buffer) (t *Tensor)

FromBuffer creates a Tensor from a backend's buffer. It requires the deviceNum information as well. The ownership of the buffer is transferred to the new Tensor.

This doesn't work for shared buffers, so it assumes the buffer is not shared.

func FromFlatDataAndDimensions

func FromFlatDataAndDimensions[T dtypes.Supported](data []T, dimensions ...int) (t *Tensor)

FromFlatDataAndDimensions creates a tensor with the given dimensions, filled with the flattened values given in `data`. The data is copied to the Tensor. The `DType` is inferred from the `data` type.

func FromScalar added in v0.11.1

func FromScalar[T dtypes.Supported](value T) (t *Tensor)

FromScalar creates a local tensor with the given scalar. The `DType` is inferred from the value.

func FromScalarAndDimensions

func FromScalarAndDimensions[T dtypes.Supported](value T, dimensions ...int) (t *Tensor)

FromScalarAndDimensions creates a local tensor with the given dimensions, filled with the given scalar value replicated everywhere. The `DType` is inferred from the value.

func FromShape

func FromShape(shape shapes.Shape) (t *Tensor)

FromShape returns a Tensor with the given shape, with the data initialized with zeros.

func FromValue

func FromValue[S MultiDimensionSlice](value S) *Tensor

FromValue returns a `Local` tensor constructed from the given multi-dimension slice (or scalar). If the rank of the `value` is larger than 1, the shape of all sub-slices must be the same.

It panics if the shape is not regular.

Notice that FromFlatDataAndDimensions is much faster if speed here is a concern.

func GobDeserialize

func GobDeserialize(decoder *gob.Decoder) (t *Tensor, err error)

GobDeserialize a Tensor from the reader.

If the tensor is only going to be directly consumed by the execution of a graph (a ML model), use GoDeserializeOnDevice instead, it works faster for some backends.

func GobDeserializeToDevice added in v0.16.0

func GobDeserializeToDevice(decoder *gob.Decoder, backend backends.Backend, deviceNums ...backends.DeviceNum) (t *Tensor, err error)

GobDeserializeToDevice deserialize a Tensor from the reader directly to on-device memory. If the tensor is expected to be consumed by graph execution, this may worked faster by avoiding an unnecessary copy.

deviceNums is optional, and for now at most one deviceNum is supported.

Returns new Tensor (with shared or onDevice storage or an error).

func Load

func Load(filePath string) (t *Tensor, err error)

Load a Local tensor from the file path given.

func (*Tensor) AssertValid

func (t *Tensor) AssertValid()

AssertValid panics if local is nil, or if its shape is invalid.

func (*Tensor) Buffer

func (t *Tensor) Buffer(backend backends.Backend, deviceNum ...backends.DeviceNum) backends.Buffer

Buffer returns the backend buffer for the tensor. It triggers the transfer from local to the backend device if the tensor is not already stored on the device.

The deviceNum is optional. But only one can be given. The default value is 0.

Careful not to finalize the tensor while the buffer is in use -- e.g.: during the execution that uses the buffer as input.

func (*Tensor) Clone added in v0.18.1

func (t *Tensor) Clone() *Tensor

Clone creates a clone of the Tensor value with shared backing with the backend -- if it supports -- or it falls back to LocalClone.

If you are trying to clone a tensor to use on a different backend, use OnDeviceClone or even LocalClone instead.

func (*Tensor) ConstBytes

func (t *Tensor) ConstBytes(accessFn func(data []byte))

ConstBytes calls accessFn with the data as a bytes slice. Even scalar values have a bytes data representation of one element. It locks the Tensor until accessFn returns.

This provides accessFn with the actual Tensor data (not a copy), and it's owned by the Tensor, and it should not be changed -- the contents of the corresponding "on device" tensors would go out-of-sync. See Tensor.MutableBytes to access a mutable version of the data as bytes.

It panics if the tensor is in an invalid state (if it was finalized), or if it is a tuple.

func (*Tensor) ConstFlatData

func (t *Tensor) ConstFlatData(accessFn func(flat any))

ConstFlatData calls accessFn with the flattened data as a slice of the Go type corresponding to the DType type. Even scalar values have a flattened data representation of one element. It locks the Tensor until accessFn returns.

It triggers a synchronous transfer from device to local, if the tensor is only on device.

This provides accessFn with the actual Tensor data (not a copy), and it's owned by the Tensor, but it should not be changed -- the contents of the corresponding "on device" tensors would go out-of-sync. See Tensor.MutableFlatData to access a mutable version of the flat data.

See Tensor.Size for the number of elements, and Tensor.LayoutStrides to calculate the offset of individual positions, given the indices at each axis.

Even scalar values have a flattened data representation of one element.

It panics if the tensor is in an invalid state (if it was finalized), or if it is a tuple.

func (*Tensor) CopyFrom added in v0.16.0

func (t *Tensor) CopyFrom(tFrom *Tensor)

CopyFrom will copy the contents from tFrom. The tensors t and tFrom must have the same shape.

This is efficient if tFrom is on-device only, in which case the device values are materialized locally into t, the receiving tensor.

func (*Tensor) DType

func (t *Tensor) DType() dtypes.DType

DType returns the DType of the tensor's shape. It is a shortcut to `Tensor.Shape().DType`.

func (*Tensor) DonateBuffer

func (t *Tensor) DonateBuffer(backend backends.Backend, deviceNum ...backends.DeviceNum) backends.Buffer

DonateBuffer returns the backend buffer for the tensor and transfers the ownership of the buffer to the caller. This may invalidate the tensor if there is no other on-device storage or local storage -- in particular this is true if using "shared buffers" (generally true for CPU-based plugins).

Mostly used internally -- by graph.Graph.Run and graph.Exec when the value in the buffer is no longer needed after execution.

It will panic if the buffer is shared (see Tensor.IsShared): shared buffers cannot be donated.

It triggers the transfer from local to the backend device if the tensor is not already stored on the device.

It doesn't finalize(release) the local tensor value.

The deviceNum is optional. But only one can be given. The default value is 0.

func (*Tensor) Equal

func (t *Tensor) Equal(otherTensor *Tensor) bool

Equal checks weather t == otherTensor. If they are the same pointer they are considered equal. If the shapes are different it returns false. If either are invalid (nil) it panics.

Slow implementation: fine for small tensors, but write something specialized for the DType if speed is desired.

func (*Tensor) FinalizeAll

func (t *Tensor) FinalizeAll()

FinalizeAll immediately frees all associated data and leave Tensor in an invalid state. Shape is cleared also.

It's the caller responsibility to ensure the tensor buffers are not being used elsewhere (like in the middle of an execution).

func (*Tensor) FinalizeLocal

func (t *Tensor) FinalizeLocal()

FinalizeLocal immediately frees the local storage copy of the tensor. If there are no on-device copies of the tensor, it becomes invalid.

If the storage is shared (see Tensor.IsShared), this is a no-op.

func (*Tensor) GoStr

func (t *Tensor) GoStr() string

GoStr converts to string, using a Go-syntax representation that can be copied&pasted back to code.

func (*Tensor) GobSerialize

func (t *Tensor) GobSerialize(encoder *gob.Encoder) (err error)

GobSerialize Tensor in binary format.

It triggers a synchronous transfer from device to local, if the tensor is only on device.

It returns an error for I/O errors. It panics for invalid tensors.

func (*Tensor) HasLocal

func (t *Tensor) HasLocal() bool

HasLocal returns whether there is an up-to-date copy of the Tensor on local storage. If false, any access to the data (e.g.: Tensor.ConstFlatData) will require a transfer (Tensor.MaterializeToLocal).

func (*Tensor) InDelta

func (t *Tensor) InDelta(otherTensor *Tensor, delta float64) bool

InDelta checks weather Abs(t - otherTensor) < delta for every element. If they are the same pointer they are considered equal. If the shapes are different it returns false. If either are invalid (nil) it panics. If the DType is not a float or complex, it also panics.

Slow implementation: fine for small tensors, but write something specialized for the DType if speed is desired.

func (*Tensor) InvalidateOnDevice

func (t *Tensor) InvalidateOnDevice()

InvalidateOnDevice destroys all on-device copies of the Tensor, so the local copy becomes the source of truth. It does nothing if the tensor is shared -- see the Tensor.IsShared method.

It's the caller's responsibility to ensure this buffer is not being used elsewhere (like in the middle of an execution).

This is automatically called if the Tensor is mutated (e.g.: Tensor.MutableFlatData) or when the on-device value is donated to the execution of a graph.

If there is no local copy of the Tensor, this will invalidate the whole tensor.

Usually, this is called automatically. Mostly for internal use.

func (*Tensor) IsLocal

func (t *Tensor) IsLocal() bool

IsLocal returns true if there is a local storage copy of the tensor. If tensor is shared (see Tensor.IsShared), it also returns true.

See MaterializeLocal to trigger a transfer/copy to the local storage.

func (*Tensor) IsOnDevice

func (t *Tensor) IsOnDevice(deviceNum backends.DeviceNum) bool

IsOnDevice checks whether the Tensor has an on-device copy on the given deviceNum.

See MaterializeOnDevices to trigger a transfer/copy to the given device.

func (*Tensor) IsScalar

func (t *Tensor) IsScalar() bool

IsScalar returns whether the tensor represents a scalar value. It is a shortcut to `Tensor.Shape().IsScalar()`.

func (*Tensor) IsShared added in v0.16.0

func (t *Tensor) IsShared() bool

IsShared returns whether the underlying tensor storage is shared with the backend engine.

In most cases, an end-user doesn't need to use this.

If true, one shouldn't access it (ConstFlatData or MutableFlatData) during the execution of a computation graph that uses it.

The Tensor implementation will try to use shared tensors where possible, since they save an extra copy.

func (*Tensor) LayoutStrides

func (t *Tensor) LayoutStrides() (strides []int)

LayoutStrides return the strides for each axis. This can be handy when manipulating the flat data.

func (*Tensor) LocalClone

func (t *Tensor) LocalClone() *Tensor

LocalClone creates a clone of the Tensor value with local backing. It will trigger a transfer from on-device data to local, if the value is not present in local memory yet.

func (*Tensor) MaterializeLocal

func (t *Tensor) MaterializeLocal()

MaterializeLocal will make sure there is a local storage of the tensor. If there isn't already a local copy, this triggers a transfer from an on-device storage to a local copy.

It's a "no-op" if using shared buffers with the accelerator. If you want to force a local copy, use LocalClone instead.

func (*Tensor) MaterializeOnDevices

func (t *Tensor) MaterializeOnDevices(backend backends.Backend, share bool, deviceNums ...backends.DeviceNum)

MaterializeOnDevices will transfer a Tensor from local storage to the given devices, if needed. Generally the user doesn't need to call this function, it's called by the libraries executing GoMLX computations automatically when needed.

If share is true, and if the backend allows for shared buffer, this will create a shared buffer, which is more economic.

- If an updated copy of the Tensor is already on the device(s), this is a no-op. - If the Tensor has already been used with a different client, this panics: one cannot mix clients on the same Tensor. - If no deviceNum is given, 0 is assumed, the default device for the client.

TODO: For now this only transfers from local storage to on-device. Implement cross-device copy on backends.

func (*Tensor) Memory

func (t *Tensor) Memory() uintptr

Memory returns the number of bytes used to store the tensor. An alias to Tensor.Shape().Memory().

func (*Tensor) MutableBytes

func (t *Tensor) MutableBytes(accessFn func(data []byte))

MutableBytes gives mutable access to the local storage of the values for the tensor. It's similar to MutableFlatData, but provide a bytes view to the same data.

If tensor is not shared, it triggers a synchronous transfer from device to local, if the tensor is only on device, and it invalidates the device storage, since it's assumed they will be out-of-date.

This returns the actual Tensor data (not a copy), and the bytes slice is owned by the Tensor -- but it's contents can be changed while inside accessFn.

See Tensor.ConstBytes for constant access to the data as bytes -- that doesn't invalidate the device storage.

func (*Tensor) MutableFlatData

func (t *Tensor) MutableFlatData(accessFn func(flat any))

MutableFlatData calls accessFn with a flat slice pointing to the Tensor data. The type of the slice is corresponds to the DType of the tensor. The contents of the slice itself can be changed until accessFn returns. During this time the Tensor is locked.

If the data is not shared with the backend (usually, only available for CPU), it invalidates and frees any copy of the data on device (e.g: GPU). It also triggers a synchronous transfer from device to local, if the tensor is only on device and not shared.

Even scalar values have a flattened data representation of one element.

This returns the actual Tensor data (not a copy), and the slice is owned by the Tensor -- but it's contents can be changed while inside accessFn.

See Tensor.ConstFlatData to access a mutable version of the flat data.

See Tensor.Size for the number of elements, and Tensor.LayoutStrides to calculate the offset of individual positions, given the indices at each axis.

It panics if the tensor is in an invalid state (if it was finalized), or if it is a tuple.

func (*Tensor) Ok

func (t *Tensor) Ok() bool

Ok returns whether the Tensor is in a valid state: it is not nil, and it hasn't been finalized.

func (*Tensor) OnDeviceClone added in v0.18.1

func (t *Tensor) OnDeviceClone(backend backends.Backend, deviceNums ...backends.DeviceNum) *Tensor

OnDeviceClone creates a clone of the tensor t that has backend storage. It also works to copy tensors to a different backend.

func (*Tensor) Rank

func (t *Tensor) Rank() int

Rank returns the rank of the tensor's shape. It is a shortcut to `Tensor.Shape().Rank()`.

func (*Tensor) Save

func (t *Tensor) Save(filePath string) (err error)

Save the Local tensor to the given file path.

It returns an error for I/O errors. It may panic if the tensor is invalid (`nil` or already finalized).

func (*Tensor) Shape

func (t *Tensor) Shape() shapes.Shape

Shape of Local, includes DType.

func (*Tensor) Size

func (t *Tensor) Size() int

Size returns the number of elements in the tensor. It is a shortcut to `Tensor.Shape().IsScalar()`.

func (*Tensor) String

func (t *Tensor) String() string

String converts to string, if not too large. It uses t.Summary(precision=4)

func (*Tensor) Summary added in v0.15.0

func (t *Tensor) Summary(precision int) string

Summary returns a multi-line summary of the contents of the Tensor. Inspired by numpy output.

func (*Tensor) ToLocal added in v0.20.0

func (t *Tensor) ToLocal() *Tensor

ToLocal forces the tensor to move its data to local (host CPU) storage and detaches itself from the backend. It returns itself to allow for cascading calls.

This is useful if using tensors across multiple backends.

If the tensor already has a local storage, there is no copy involved. Any on-device storage is freed.

func (*Tensor) Value

func (t *Tensor) Value() any

Value returns a multidimensional slice (except if shape is a scalar) containing a copy of the values stored in the tensor. This is expensive, and usually only used for smaller tensors in tests and to print results.

If the local tensor is empty it panics with the corresponding error.

Directories

Path Synopsis
Package images provides several functions to transform images back and forth from tensors.
Package images provides several functions to transform images back and forth from tensors.
Package numpy allows one to read/write tensors to Python's NumPy npy and npz file formats.
Package numpy allows one to read/write tensors to Python's NumPy npy and npz file formats.

Jump to

Keyboard shortcuts

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