tensor

package
v0.4.0 Latest Latest
Warning

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

 Go to latest Published: Jun 16, 2023 License: Apache-2.0 Imports: 11 Imported by: 0

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

Constants

This section is empty.

Variables

View Source 
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 

func FlatFromRef[T shapes.Number](ref *LocalRef) []T

FlatFromRef returns the flattened data acquired by `ref` as a slice of the corresponding DType type. It is the generics version of LocalRef.Flat()

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 returns nil.

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.

To create it, either create a Local tensor first and then convert it to Device, or use the output of the execution of a computation graph -- they return Device tensors.

It implements the Tensor interface.

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 other devices and the local tensor (if any), one can use FinalizeAll().

func InternalNewDevice added in v0.2.1 

func InternalNewDevice(buffer *xla.OnDeviceBuffer) (deviceT *Device)

InternalNewDevice creates a Device tensor from XLA's OnDeviceBuffer structure. Internal implementation, most users don't need to use this.

func MakeDeviceWithError added in v0.2.1 

func MakeDeviceWithError(err error) *Device

MakeDeviceWithError creates a device tensor with the given error.

func (Device) AddDevice 

func (c Device) AddDevice(device *Device) *Device

AddDevice to the internal cache, and returns itself for convenience.

func (Device) AddLocal 

func (c Device) AddLocal(local *Local) *Local

AddLocal to cache and returns the local tensor for convenience.

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) DType 

func (device *Device) DType() shapes.DType

DType returns the DType of the tensor's shape.

func (Device) Device 

func (c Device) Device(hasClient HasClient, deviceNum int) *Device

Device either uses a cached value on device already or it transfers local data to the shapedBuffer store of values (OnDeviceBuffer) and returns a tensor.Device reference -- value is cached for future calls. This is used for instance to transfer parameters when executing a graph.

func (*Device) Empty 

func (device *Device) Empty() bool

Empty returns whether Local is holding no data or is in an error state. It's similar to a "nil" state for Local.

func (*Device) Error 

func (device *Device) Error() error

Error returns the message that caused an error state.

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 scarse, this allows for finer control of memory usage.

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) IsTuple 

func (device *Device) IsTuple() bool

IsTuple returns whether Local is a tuple.

func (Device) Local 

func (c Device) Local() *Local

Local will transfer data from the Device storage to a Local tensor.

func (*Device) Ok 

func (device *Device) Ok() bool

Ok returns whether the shapedBuffer is not empty and has no error.

func (*Device) Rank 

func (device *Device) Rank() int

Rank returns the rank fo the tensor's shape.

func (*Device) Shape 

func (device *Device) Shape() shapes.Shape

Shape returns the shape of the Device.

func (*Device) ShapedBuffer 

func (device *Device) ShapedBuffer() *xla.OnDeviceBuffer

ShapedBuffer returns the underlying XLA structure. Internal usage only.

func (*Device) SplitTuple 

func (device *Device) SplitTuple() []*Device

SplitTuple splits a device tensor into its elements. In case of error, returns nil.

This makes the current device tensor invalid.

func (*Device) SplitTupleError 

func (device *Device) SplitTupleError() ([]*Device, error)

SplitTupleError splits a device tensor into its elements. In case of error, return the error.

This makes the current device tensor invalid.

func (*Device) String 

func (device *Device) String() string

String converts to string, by converting (transferring) the tensor to local and then using Local.String().

func (*Device) Value 

func (device *Device) Value() any

Value returns a multidimensional slice (except if shape is a scalar) containing a copy of the tensor values. It is just a shortcut to calling `device.Local().Value()`. See Local and Local.Value for details.

type HasClient 

type HasClient interface {
	Client() *xla.Client
}

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's 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 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.Flat below) with the flat data, and then creates the sub-slices point to the middle of it.
  • `Local.Flat() 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.CopyFlat(dst any) error`: similar to Local.Flat, but instead of creating a new slice, it copies to an already provided slice. Useful when looping over results. It returns an error if `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 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 Ok, Empty and Error.

func FromFlatDataAndDimensions added in v0.2.1 

func FromFlatDataAndDimensions[T shapes.Supported](data []T, dimensions ...int) (local *Local)

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

func FromScalarAndDimensions added in v0.2.1 

func FromScalarAndDimensions[T shapes.Supported](value T, dimensions ...int) (local *Local)

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) (local *Local)

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 rank of the `value` is larger than 1, the shape of all sub-slices must be the same.

It returns a tensor.Local with error if shape is not regular.

func GobDeserialize added in v0.2.0 

func GobDeserialize(decoder *gob.Decoder) (local *Local, err error)

GobDeserialize a Tensor from the reader. Returns new tensor.Local or an error.

func Load added in v0.3.0 

func Load(filePath string) (local *Local, err error)

Load a Local tensor from the file path given.

func MakeLocalTuple 

func MakeLocalTuple(tensors ...*Local) *Local

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 

func MakeLocalTupleAny(values ...any) *Local

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 MakeLocalWithError added in v0.2.1 

func MakeLocalWithError(err error) *Local

MakeLocalWithError creates a local tensor with the given error.

func (*Local) AcquireData added in v0.2.1 

func (local *Local) AcquireData() *LocalRef

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 returns nil if Local tensor is in an invalid state, or if it is a tuple.

func (Local) AddDevice 

func (c Local) AddDevice(device *Device) *Device

AddDevice to the internal cache, and returns itself for convenience.

func (Local) AddLocal 

func (c Local) AddLocal(local *Local) *Local

AddLocal to cache and returns the local tensor for convenience.

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) CopyFlat added in v0.2.1 

func (local *Local) CopyFlat(dst any) error

CopyFlat 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.

func (*Local) DType 

func (local *Local) DType() shapes.DType

DType returns the DType of the tensor's shape. Shortcut to `local.Shape().DType`

func (*Local) Data 

func (local *Local) Data() any

Data returns a slice with the consecutive data of the corresponding DType type. Consider the generic function Data[L]() if you know the type upfront.

It returns nil is Local tensor is in an invalid state, or if it is a tuple.

func (Local) Device 

func (c Local) Device(hasClient HasClient, deviceNum int) *Device

Device either uses a cached value on device already or it transfers local data to the shapedBuffer store of values (OnDeviceBuffer) and returns a tensor.Device reference -- value is cached for future calls. This is used for instance to transfer parameters when executing a graph.

func (*Local) Empty 

func (local *Local) Empty() bool

Empty returns whether Local is holding no data. It's similar to a "nil" state for Local.

func (*Local) Error 

func (local *Local) Error() error

Error returns the message that caused an error state.

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 on 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) Flat added in v0.2.1 

func (local *Local) Flat() any

Flat returns a copy of the flattened contents of the Local tensor as a slice of the type matching the tensor's DType (see shapes.TypeForDType).

See Local.LayoutStrides to calculate the offset of individual positions.

It the tensor is a scalar, it still returns a slice with one element.

If tensor is empty, in error of if it is a tuple, returns nil.

func (*Local) GoStr 

func (local *Local) GoStr() string

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 

func (local *Local) GobSerialize(encoder *gob.Encoder) (err error)

GobSerialize Local tensor in binary format.

func (*Local) IsTuple 

func (local *Local) IsTuple() bool

IsTuple returns whether Local is a tuple.

func (*Local) LayoutStrides added in v0.2.1 

func (local *Local) LayoutStrides() (strides []int)

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

func (*Local) Literal 

func (local *Local) Literal() *xla.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) Local 

func (c Local) Local() *Local

Local will transfer data from the Device storage to a Local tensor.

func (*Local) Ok 

func (local *Local) Ok() bool

Ok returns whether local is both not empty and is not in error.

func (*Local) Rank 

func (local *Local) Rank() int

Rank returns the rank fo the tensor's shape. Shortcut to `local.Shape().Rank()`

func (*Local) Save added in v0.3.0 

func (local *Local) Save(filePath string) (err error)

Save the Local tensor to the given file path.

func (*Local) Shape 

func (local *Local) Shape() shapes.Shape

Shape of Local, includes DType.

func (*Local) SplitTuple 

func (local *Local) SplitTuple() (tensors []*Local, err error)

SplitTuple splits the Tuple tensor into its components. This unfortunately destroys the current Local, emptying it.

func (*Local) String 

func (local *Local) String() string

String converts to string, if not too large.

func (*Local) StringN 

func (local *Local) StringN(n int) string

StringN converts to string, displaying at most n elements. TODO: nice pretty-print version, even for large tensors.

func (*Local) Value 

func (local *Local) Value() any

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 a version the one can mutate.

If tensor is empty, in error of if it is a tuple, returns nil.

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) Bytes added in v0.2.1 

func (ref *LocalRef) Bytes() []byte

Bytes returns the same memory as Flat, but the raw slice of bytes, with the proper size in bytes.

It is only valid while LocalRef hasn't been released.

func (*LocalRef) Flat added in v0.2.1 

func (ref *LocalRef) Flat() any

Flat returns the flattened data as a slice of the corresponding DType type.

See Local.LayoutStrides to calculate the offset of individual positions.

It is only valid while LocalRef hasn't been released.

It returns nil is Local tensor is in an invalid state, or if it is a tuple.

func (*LocalRef) Ok added in v0.2.1 

func (ref *LocalRef) Ok() bool

Ok returns whether ref, and the Local tensor it points to are valid.

func (*LocalRef) Release added in v0.2.1 

func (ref *LocalRef) Release()

Release acquired data reference to a Local tensor. If LocalRef is nil (in case a call to Local.AcquireData failed), it is a no-op.

type Tensor 

type Tensor interface {
	// Local version of the tensor. If the underlying tensor is Local already, it's a no-op. Otherwise, the tensor
	// contents are transferred locally. It uses a cache system, so if tensor was already local no transfer happens.
	Local() *Local

	// Device version of the tensor. If the underlying tensor is on the given Device already, it's a no-op. Otherwise,
	// the tensor contents are transferred to the device. It uses a cache system, so if tensor was already local no
	// transfer happens.
	Device(client HasClient, deviceNum int) *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

	// Error returns the message that caused an error state.
	Error() error

	// Ok returns whether the tensor is in an invalid state or empty.
	Ok() bool

	// FinalizeAll immediately frees the data from all versions of the Tensor -- Local or on Device, and make the
	// tensor invalid. This calls Finalize on the cached Local and all Device tensors.
	FinalizeAll()
}

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 

func FromAnyValue(value any) Tensor

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 returns a tensor. with an error (see `Tensor.Error()`) if `value` type is unsupported or the shape is not regular.

Source Files

View all Source files

Directories

Path Synopsis
Package image provides several functions to transform images back and forth from tensors.
 Package image provides several functions to transform images back and forth from tensors.

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.