tensor

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: 14 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.Supported](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() 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) 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) 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) 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 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 

func (device *Device) IsFinalized() bool

IsFinalized returns true if the tensor has already been "finalized", and its data freed. It implements Tensor.IsFinalized.

func (*Device) IsTuple 

func (device *Device) IsTuple() bool

IsTuple returns whether Local is a tuple.

func (*Device) Local 

func (device *Device) Local() *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) Rank 

func (device *Device) Rank() int

Rank returns the rank of 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.

This makes the current device tensor invalid -- but not any associated Local (or other Device) tensors.

func (*Device) String 

func (device *Device) String() string

String converts to string, by converting (transferring) the tensor to local and then using Local.String(). If the tensor is larger than MaxStringSize, it doesn't convert the tensor to local, and instead only prints its shape (with no transfer cost).

func (Device) Value 

func (c 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 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 

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 `data` type.

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

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 (*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 panics 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) 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 

func (local *Local) CopyData(dst any)

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 

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

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

func (Local) Device 

func (c Local) Device(hasClient HasClient, deviceNum int) *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 

func (local *Local) FlatCopy() any

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 

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.

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

func (*Local) IsFinalized added in v0.5.0 

func (local *Local) IsFinalized() bool

IsFinalized returns true if the tensor has already been "finalized", and its data freed. It implements Tensor.IsFinalized.

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

Local implements Tensor. It simply returns itself.

func (*Local) Rank 

func (local *Local) Rank() int

Rank returns the rank of 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.

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

func (*Local) Shape 

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

Shape of Local, includes DType.

func (*Local) SplitTuple 

func (local *Local) SplitTuple() (tensors []*Local)

SplitTuple splits the Tuple tensor into its components. This destroys the local tensor, making it invalid.

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

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 `ref` 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. 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.

func (*LocalRef) Release added in v0.2.1 

func (ref *LocalRef) Release()

Release acquired data reference to a Local tensor. If `ref` 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.
	// 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 

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 panics with an 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.