Documentation ¶
Overview ¶
Package tensor provides a `Tensor` interface with 2 different implementations: `Local` and `Device`, they differ on where their values are stored: in the local (host) CPU, or on an accelerator device (TPU, GPU, but could be also the CPU if no accelerators are available).
The main use of tensors are to be used as input and output of GoMLX computation graph.
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.
This implementation uses `gomlx.types.Shape` to represent the shape, and (for now) only explicitly supports dense representations of the data. The underlying implementation of both `Local` and `Device` implementation are wrappers to similar XLA data representations.
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. Because the tensors are immutable, the transferring is cached, so if `Tensor.Local()` or `Tensor.Device()` is called multiple times, the price is paid only once.
To facilitate managing this, this package also implements a cache system whereas a Local or Device tensor caches references to their counterparts -- notice there may be multiple Device tensors (one for each device). This is all exposed through the Tensor interface.
Concurrency: the cache system is safe from concurrency point of view. But management of the conflicting uses of the content of the tensors themselves is left for the users -- TODO: this needs to be improved.
See details on the documentation of Tensor, Device and Local structures.
Index ¶
- Variables
- func FlatFromRef[T shapes.Supported](ref *LocalRef) []T
- type Device
- func (c Device) AddDevice(device *Device) *Device
- func (c Device) AddLocal(local *Local) *Local
- func (device *Device) AssertValid()
- func (device *Device) ClearCache()
- func (c Device) ClearDevice(device *Device)
- func (c Device) ClearLocal()
- func (c Device) CurrentDevice() *Device
- func (device *Device) DType() shapes.DType
- func (c Device) Device(hasClient HasClient, deviceNum int) *Device
- func (device *Device) Finalize()
- func (device *Device) FinalizeAll()
- func (device *Device) IsFinalized() bool
- func (device *Device) IsTuple() bool
- func (device *Device) Local() *Local
- func (device *Device) Rank() int
- func (device *Device) Shape() shapes.Shape
- func (device *Device) ShapedBuffer() *xla.OnDeviceBuffer
- func (device *Device) SplitTuple() []*Device
- func (device *Device) String() string
- func (c Device) Value() any
- type HasClient
- type Local
- func FromFlatDataAndDimensions[T shapes.Supported](data []T, dimensions ...int) (local *Local)
- func FromScalarAndDimensions[T shapes.Supported](value T, dimensions ...int) (local *Local)
- func FromShape(shape shapes.Shape) (local *Local)
- func FromValue[S shapes.MultiDimensionSlice](value S) *Local
- func GobDeserialize(decoder *gob.Decoder) (local *Local, err error)
- func Load(filePath string) (local *Local, err error)
- func MakeLocalTuple(tensors ...*Local) *Local
- func MakeLocalTupleAny(values ...any) *Local
- func (local *Local) AcquireData() *LocalRef
- func (c Local) AddDevice(device *Device) *Device
- func (c Local) AddLocal(local *Local) *Local
- func (local *Local) AssertValid()
- func (local *Local) AssertValidAndNoTuple()
- func (local *Local) ClearCache()
- func (c Local) ClearDevice(device *Device)
- func (c Local) ClearLocal()
- func (local *Local) CopyData(dst any)
- func (c Local) CurrentDevice() *Device
- func (local *Local) DType() shapes.DType
- func (c Local) Device(hasClient HasClient, deviceNum int) *Device
- func (local *Local) Finalize()
- func (local *Local) FinalizeAll()
- func (local *Local) FlatCopy() any
- func (local *Local) GoStr() string
- func (local *Local) GobSerialize(encoder *gob.Encoder) (err error)
- func (local *Local) IsFinalized() bool
- func (local *Local) IsTuple() bool
- func (local *Local) LayoutStrides() (strides []int)
- func (local *Local) Literal() *xla.Literal
- func (local *Local) Local() *Local
- func (local *Local) Rank() int
- func (local *Local) Save(filePath string) (err error)
- func (local *Local) Shape() shapes.Shape
- func (local *Local) SplitTuple() (tensors []*Local)
- func (local *Local) String() string
- func (local *Local) StringN(n int) string
- func (local *Local) Value() any
- type LocalRef
- type Tensor
Constants ¶
This section is empty.
Variables ¶
var MaxStringSize = 500
MaxStringSize is the largest Local tensor that is actually returned by String() is requested.
Functions ¶
func FlatFromRef ¶ added in v0.2.1
FlatFromRef returns the flattened data acquired by `ref` as a slice of the corresponding DType type. It is the "generics" version of LocalRef.Flat() This is not a copy, but a pointer to the underlying data, that can be changed.
See Local.LayoutStrides to calculate the offset of individual positions.
If the DType of the tensor is not compatible with the requested number type, or if tensor is a tuple, it throws a panic with an error message.
Types ¶
type Device ¶
type Device struct {
// contains filtered or unexported fields
}
Device represents a tensor (or tuple) stored on a "device": typically an accelerator (GPU, TPU), but it can also be the normal memory if using the CPU. The object doesn't offer much functionality, it's simply used as input and output of graph execution. A Device is represented by a `ClientId`, typically provided by the `graph.Manager` object, and a device number -- in case there are multiple GPUs/TPUs.
To access it, one has to convert it to a `Local` tensor, which some methods (like `Value()`) do automatically. But converting to `Local` usually involves copying across devices, so avoid it when not needed -- for instance, variables (or weights) of models typically don't need to be copied to local, they can simply be left on-device until they need to be saved for instance.
To create it, either create a Local tensor first and then convert it to Device (see `Tensor.Device() method`), or use the output of a computation graph execution (see `graph` package) -- they return Device tensors.
It implements the Tensor interface.
When dealing with large tensors, one will want to carefully manage its life cycle. `Device` provides a method called `Finalize()` to immediately release its data memory (managed in XLA C++) -- it is also called automatically during garbage collection. And to release all associated versions of the tensor, including the copies on other devices and the Local tensor (if any), one can use `FinalizeAll()`.
func InternalNewDevice ¶ added in v0.2.1
func InternalNewDevice(buffer *xla.OnDeviceBuffer) (device *Device)
InternalNewDevice creates a Device tensor from XLA's OnDeviceBuffer structure.
Internal implementation, most users shouldn't use this. Instead, one creates a `Local` and converts it to a `Device` using the `Tensor.Device()` method.
func (*Device) AssertValid ¶ added in v0.5.0
func (device *Device) AssertValid()
AssertValid panics if device is nil, if its shape is invalid or if it has already been finalized (freed).
func (*Device) ClearCache ¶
func (device *Device) ClearCache()
ClearCache disconnects the device tensor to any corresponding local data. Internal usage only.
func (Device) ClearDevice ¶
func (c Device) ClearDevice(device *Device)
ClearDevice from cache, and leaves the device tensor passed without a cache.
func (Device) ClearLocal ¶
func (c Device) ClearLocal()
ClearLocal cache, and leaves the local cached tensor without a cache.
func (Device) CurrentDevice ¶ added in v0.4.1
func (c Device) CurrentDevice() *Device
CurrentDevice returns the current Device tensor backing this tensor, if there is any. If the tensor is Local, this returns nil.
If there is more than one Device tensor, this returns the first one.
func (Device) Device ¶
Device implements Tensor.Device. It returns a `Device` version of the tensor. If the underlying tensor is on the given `Device` already, it's a no-op. If there is already a cache of a corresponding `Device` tensor on the given device, that is returned. Otherwise, the contents of the `Local` tensor (if available), or one of the on-device tensors are transferred to a `Device` tensor on the given `deviceNum`.
func (*Device) Finalize ¶
func (device *Device) Finalize()
Finalize releases the memory associated with the Device tensor, it becomes empty.
This is called automatically when garbage-collected. But since sometimes device memory is scarce, this allows for finer control of memory usage.
This can be called more than once: after the first time it doesn't do anything, since the data has already been released.
func (*Device) FinalizeAll ¶
func (device *Device) FinalizeAll()
FinalizeAll releases the memory associated with all copies of the tensor, local and on device(s)). And then mark them as empty.
func (*Device) IsFinalized ¶ added in v0.5.0
IsFinalized returns true if the tensor has already been "finalized", and its data freed. It implements Tensor.IsFinalized.
func (*Device) Local ¶
Local will transfer data from the Device storage to a Local tensor. If the tensor has already been converted, return the associated cached copy.
func (*Device) ShapedBuffer ¶
func (device *Device) ShapedBuffer() *xla.OnDeviceBuffer
ShapedBuffer returns the underlying XLA structure. Internal usage only.
func (*Device) SplitTuple ¶
SplitTuple splits a device tensor into its elements.
This makes the current device tensor invalid -- but not any associated Local (or other Device) tensors.
type HasClient ¶
HasClient accepts anything that can return a xla.Client. That includes xla.Client itself and graph.Manager.
type Local ¶
type Local struct {
// contains filtered or unexported fields
}
Local tensor represents a multidimensional array of one of the supported types (see shapes.Number) that is stored "locally" (in CPU memory), albeit under XLA (that is C++) management. It is meant to be used means of feeding and retrieving data for GoMLX computation graphs.
It is shape can be from a scalar to an arbitrary rank (number of dimensions). See `Shape()` to get information about its dimensions and the underlying `DType`.
As a special case, it can also hold a tuple of Local tensors (recursive definition).
It implements the generic `Tensor` interface, and it provides some specialized functionality that assumes the data is local (in CPU). The Tensor interface includes conversion back-and-forth to a Device tensor. The converted tensors are cached, so there is no extra cost in doing it multiple times.
When dealing with large tensors, one will want to carefully manage its life cycle. It provides a method called Finalize() to immediately release its data memory (managed in XLA C++) -- it is also called automatically during garbage collection. And to release all associated versions of the tensor, including the copies on-device, there is FinalizeAll().
There are various ways to construct a Local tensor:
- `FromValue[S shapes.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. E.g.: `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, it is itself returned.
- `FromShape(shape shapes.Shape)`: creates a `Local` tensor with the given shape, and uninitialized values. See below how to mutate its values with a `LocalRef`.
- `FromScalarAndDimensions[T shapes.Supported](value T, dimensions ...int)`: creates a Local 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 Local tensor with the given dimensions, and set the flattened values with the given data. `T` must be one of the supported types.
There are various ways to access the contents of a Local tensor:
- `Local.Value() any`: it creates a copy of the tensor contents to a Go type. A scalar if the underlying shape is a scalar or a (multidimensional) slice. If returning a multidimensional slice, it first creates a flat slice (see Local.FlatCopy below) with the flat data, and then creates the sub-slices point to the middle of it.
- `Local.FlatCopy() any`: similar to Value() it creates a copy of the flattened (only one dimension) contents to a Go slice -- even if it is a scalar, it will return a slice with one value.
- `Local.CopyData(dst any) error`: similar to Local.FlatCopy, but instead of creating a new slice, it copies to an already provided slice. Useful when looping over results. It returns an error if the `dst` type is not compatible or not if there is not the correct space in `dst`, that is, if `len(dst) != Local.Shape().Size()`.
The methods above are all for reading. To mutate the contents of the Local tensor after it was created, or to access it directly without any copies -- relevant if using large tensors -- one has to "acquire" the data. That makes sure the GC is not going to collect the data (since it's in C++, Go doesn't know about it). Call `Local.AcquireData()` to get a `LocalRef`. Then do something like `defer LocalRef.Release()` to make sure the data is released after its used (otherwise it will never be freed). And with the LocalRef one can use:
- `LocalRef.Flat() any`: like `Local.Flat` returns a slice with the flattened data, but it points to the underlying c++ flattened data. Mutating this slice contents changes the tensor content directly. The returned slice is only valid while LocalRef is not released (LocalRef.Release), and the tensor is not finalized (Local.Finalize).
- `FlatFromRef[T shapes.Number](ref *LocalRef) []T`: a generic function version of `LocalRef.Flat()` (it does a couple fewer runtime casting in the process).
- `LocalRef.Bytes() []byte` returns a slice of bytes pointing directly to the storage bytes of tensor. Convenient for saving or restoring tensor contents. See also GobSerialize and GobDeserialize.
Notice there is not a library of tensor functions (math, or otherwise) to operate on them. To do math on tensors, instead use the computation graph engine (package `graph`).
Local can be in an empty state (for instance, after calling Finalize), and in an erroneous state -- after some invalid operation. These can be checked and tested with the methods AssertValid, Empty and Error.
func FromFlatDataAndDimensions ¶ added in v0.2.1
FromFlatDataAndDimensions creates a local tensor with the given dimensions, filled with the flattened values given in `data`. The `DType` is inferred from the `data` type.
func FromScalarAndDimensions ¶ added in v0.2.1
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 ¶
FromShape creates a `Local` tensor with the given shape, with the data uninitialized. See `Local.AcquireData()` to mutate the data after the `Local` tensor is created.
func FromValue ¶
func FromValue[S shapes.MultiDimensionSlice](value S) *Local
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 ¶ added in v0.2.0
GobDeserialize a Tensor from the reader. Returns new tensor.Local or an error.
func MakeLocalTuple ¶
MakeLocalTuple compose local tensors into a Tuple. The individual tensors are destroyed in the process, as the tuple takes ownership of its parts.
func MakeLocalTupleAny ¶
MakeLocalTupleAny composes values into local tensor. Values can be any value that can be converted to a *Local tensor, or a *Local tensor. Similar to MakeLocalTuple, but more permissible.
func (*Local) AcquireData ¶ added in v0.2.1
AcquireData returns a LocalRef that can be used to access the underlying C++ data directly.
Once acquired, it has to be manually released -- it may leak tensors if not released. The recommended way it to pair it with a deferred release, as in the example below:
dataRef := local.AcquireData() defer dataRef.Release() // do something with dataRef...
It does not work for tuples: those need to be split first.
It panics if Local tensor is in an invalid state, or if it is a tuple.
func (*Local) AssertValid ¶ added in v0.5.0
func (local *Local) AssertValid()
AssertValid panics if local is nil, or if its shape is invalid.
func (*Local) AssertValidAndNoTuple ¶ added in v0.5.0
func (local *Local) AssertValidAndNoTuple()
AssertValidAndNoTuple both asserts it's a valid tensor and that it's not a tuple. It panics with corresponding errors if violated.
func (*Local) ClearCache ¶
func (local *Local) ClearCache()
ClearCache disconnect the local tensor to the corresponding cache data, which holds the pointers to the on Device versions of the tensor. This should be called if the `Local` tensor contents are mutated, but its cache is pointing to a lagging version of a Device tensor. See discussion in Tensor interface.
func (Local) ClearDevice ¶
func (c Local) ClearDevice(device *Device)
ClearDevice from cache, and leaves the device tensor passed without a cache.
func (Local) ClearLocal ¶
func (c Local) ClearLocal()
ClearLocal cache, and leaves the local cached tensor without a cache.
func (*Local) CopyData ¶ added in v0.5.0
CopyData contents of the tensor to dst. The parameter `dst` must be a slice of the corresponding type that matches the tensor `DType`, see `shapes.TypeForDType`. Data is copied flat, as a 1D slice. Even for scalars, it is copied as a slice of size 1.
It panics if the tensor is empty or dst doesn't hold enough space or the right type.
func (Local) CurrentDevice ¶ added in v0.4.1
func (c Local) CurrentDevice() *Device
CurrentDevice returns the current Device tensor backing this tensor, if there is any. If the tensor is Local, this returns nil.
If there is more than one Device tensor, this returns the first one.
func (*Local) DType ¶
DType returns the DType of the tensor's shape. Shortcut to `local.Shape().DType`
func (Local) Device ¶
Device implements Tensor.Device. It returns a `Device` version of the tensor. If the underlying tensor is on the given `Device` already, it's a no-op. If there is already a cache of a corresponding `Device` tensor on the given device, that is returned. Otherwise, the contents of the `Local` tensor (if available), or one of the on-device tensors are transferred to a `Device` tensor on the given `deviceNum`.
func (*Local) Finalize ¶
func (local *Local) Finalize()
Finalize releases the memory associated with the local tensor. It becomes Empty() = true. It mutates the tensor, but it's handy in case one is dealing with large data. See discussion on storage and mutability in the package documentation.
func (*Local) FinalizeAll ¶
func (local *Local) FinalizeAll()
FinalizeAll releases the memory associated with all copies of the tensor (local and on device), and mark them as empty.
func (*Local) FlatCopy ¶ added in v0.8.0
FlatCopy returns a copy of `Local` tensor's flattened contents as a slice of the type matching the tensor's DType (see shapes.TypeForDType).
See Local.LayoutStrides to calculate the offset of individual positions.
If the tensor is a scalar, it still returns a slice with one element.
If tensor is invalid (already finalized?), of if is a tuple, it panics with an error.
func (*Local) GoStr ¶
GoStr converts to string, using a Go-syntax representation that can be copied&pasted back to code.
func (*Local) GobSerialize ¶ added in v0.2.0
GobSerialize Local tensor in binary format.
It returns an error for I/O errors. It panics for invalid tensors.
func (*Local) IsFinalized ¶ added in v0.5.0
IsFinalized returns true if the tensor has already been "finalized", and its data freed. It implements Tensor.IsFinalized.
func (*Local) LayoutStrides ¶ added in v0.2.1
LayoutStrides return the strides for each axis. This can be handy when manipulating the flat data.
func (*Local) Literal ¶
Literal returns the internal storage of the Local value. Internal usage only, used only by new Op implementations. If you need access to the underlying data, use AcquireData instead.
func (*Local) Rank ¶
Rank returns the rank of the tensor's shape. Shortcut to `local.Shape().Rank()`
func (*Local) Save ¶ added in v0.3.0
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 (*Local) SplitTuple ¶
SplitTuple splits the Tuple tensor into its components. This destroys the local tensor, making it invalid.
func (*Local) StringN ¶
StringN converts to string, displaying at most n elements. TODO: nice pretty-print version, even for large tensors.
func (*Local) Value ¶
Value returns a multidimensional slice (except if shape is a scalar) containing the values. The returned `any` value can be cast to the appropriate Go type.
It returns a copy of the underlying data. See AcquireData to access the data directly (instead of a copy), which is also mutable.
If the local tensor is empty or a tuple, it panics with the corresponding error.
type LocalRef ¶ added in v0.2.1
type LocalRef struct {
// contains filtered or unexported fields
}
LocalRef is a live-reference to a Local tensors data (managed by C++). It keeps it alive (from the GC), and should be released with Release at the end of its use -- see example in Local.AcquireData.
It provides direct access to the underlying data.
func (*LocalRef) AssertValid ¶ added in v0.5.0
func (ref *LocalRef) AssertValid()
AssertValid panics if reference is empty or has already been freed.
func (*LocalRef) Bytes ¶ added in v0.2.1
Bytes returns the same memory as Flat, but the raw slice of bytes, with the proper size in bytes.
It is only valid while `ref` hasn't been released.
func (*LocalRef) Flat ¶ added in v0.2.1
Flat returns the flattened data as a slice of the corresponding DType type. This is not a copy, but a pointer to the underlying data, that can be changed.
See Local.LayoutStrides to calculate the offset of individual positions.
It is only valid while `ref` hasn't been released.
It panics if Local tensor is in an invalid state, or if it is a tuple.
type Tensor ¶
type Tensor interface { // Local version of the tensor. // If the underlying tensor is `Local` already, it's a no-op. // If there is already a cache of a corresponding `Local` tensor, that is returned. // Otherwise, the contents of one of the on-device tensors (if there is more than one) are // transferred locally. Local() *Local // Device version of the tensor. // If the underlying tensor is on the given `Device` already, it's a no-op. // If there is already a cache of a corresponding `Device` tensor on the given device, that is returned. // Otherwise, the contents of the `Local` tensor (if available), or one of the on-device tensors are // transferred to a `Device` tensor on the given `deviceNum`. Device(client HasClient, deviceNum int) *Device // CurrentDevice returns the current Device tensor backing this tensor, if there is any. // If the tensor is Local, this returns nil. // // If there is more than one Device tensor, this returns the first one. CurrentDevice() *Device // Shape of the tensor. Shape() shapes.Shape // DType of the tensor's shape. DType() shapes.DType // Rank of the tensor's shape. Rank() int // String returns a printable version of the tensor. This may lead to a transfer from a Device tensor // with the Local(). String() string // Value returns a multidimensional slice (except if shape is a scalar) containing the values. // If the underlying tensor is on device (E.g.: GPU), it's transferred locally with Local(). Value() any // FinalizeAll immediately frees the data from all versions of the Tensor -- local or on-device, and makes the // tensor invalid. // This calls Finalize on the cached Local and all Device tensors. FinalizeAll() // IsFinalized returns true if the tensor has already been "finalized", and its // data freed. IsFinalized() bool }
Tensor represents a 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.
Tensor can be implemented by a tensor.Local or tensor.Device, which reflects whether the data is stored in the local CPU on or the device actually running the computation: an accelerator like a GPU or the CPU as well.
Local and Device tensors can be converted to each other -- there is a transferring cost to that. There is a cache system to prevent duplicate transfers, but it assumes immutability -- call Local.ClearCache after mutating a Local tensor. Device tensors are immutable.
More details in the `tensor` package documentation.
func FromAnyValue ¶
FromAnyValue is a non-generic version of FromValue that returns a tensor.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.