ML Community Day is November 9! Join us for updates from TensorFlow, JAX, and more Learn more

TensorBuffer

public abstract class TensorBuffer
Known Direct Subclasses

Represents the data buffer for either a model's input or its output.

Public Methods

static TensorBuffer
createDynamic(DataType dataType)
Creates an empty dynamic TensorBuffer with specified DataType.
static TensorBuffer
createFixedSize(int[] shape, DataType dataType)
Creates a TensorBuffer with specified shape and DataType.
static TensorBuffer
createFrom(TensorBuffer buffer, DataType dataType)
Creates a TensorBuffer deep-copying data from another, with specified DataType.
ByteBuffer
getBuffer()
Returns the data buffer.
abstract DataType
getDataType()
Returns the data type of this buffer.
int
getFlatSize()
Gets the flatSize of the buffer.
abstract float[]
getFloatArray()
Returns a float array of the values stored in this buffer.
abstract float
getFloatValue(int absIndex)
Returns a float value at a given index.
abstract int[]
getIntArray()
Returns an int array of the values stored in this buffer.
abstract int
getIntValue(int absIndex)
Returns an int value at a given index.
int[]
getShape()
Gets the current shape.
abstract int
getTypeSize()
Returns the number of bytes of a single element in the array.
boolean
isDynamic()
Returns if the TensorBuffer is dynamic sized (could resize arbitrarily).
abstract void
loadArray(int[] src, int[] shape)
Loads an int array into this buffer with specific shape.
abstract void
loadArray(float[] src, int[] shape)
Loads a float array into this buffer with specific shape.
void
loadArray(float[] src)
Loads a float array into this buffer.
void
loadArray(int[] src)
Loads an int array into this buffer.
void
loadBuffer(ByteBuffer buffer)
Loads a byte buffer into this TensorBuffer.
void
loadBuffer(ByteBuffer buffer, int[] shape)
Loads a byte buffer into this TensorBuffer with specific shape.

Inherited Methods

Public Methods

public static TensorBuffer createDynamic (DataType dataType)

Creates an empty dynamic TensorBuffer with specified DataType. The shape of the created TensorBuffer is {0}.

Dynamic TensorBuffers will reallocate memory when loading arrays or data buffers of different buffer sizes. Here are some examples:

 // Creating a float dynamic TensorBuffer:
 TensorBuffer tensorBuffer = TensorBuffer.createDynamic(DataType.FLOAT32);
 // Loading a float array:
 float[] arr1 = new float[] {1, 2, 3};
 tensorBuffer.loadArray(arr, new int[] {arr1.length});
 // loading another float array:
 float[] arr2 = new float[] {1, 2, 3, 4, 5};
 tensorBuffer.loadArray(arr, new int[] {arr2.length});
 // loading a third float array with the same size as arr2, assuming shape doesn't change:
 float[] arr3 = new float[] {5, 4, 3, 2, 1};
 tensorBuffer.loadArray(arr);
 // loading a forth float array with different size as arr3 and omitting the shape will result
 // in error:
 float[] arr4 = new float[] {3, 2, 1};
 tensorBuffer.loadArray(arr); // Error: The size of byte buffer and the shape do not match.
 

Parameters
dataType The dataType of the TensorBuffer to be created.

public static TensorBuffer createFixedSize (int[] shape, DataType dataType)

Creates a TensorBuffer with specified shape and DataType. Here are some examples:

 // Creating a float TensorBuffer with shape {2, 3}:
 int[] shape = new int[] {2, 3};
 TensorBuffer tensorBuffer = TensorBuffer.createFixedSize(shape, DataType.FLOAT32);
 
 // Creating an uint8 TensorBuffer of a scalar:
 int[] shape = new int[] {};
 TensorBuffer tensorBuffer = TensorBuffer.createFixedSize(shape, DataType.UINT8);
 
 // Creating an empty uint8 TensorBuffer:
 int[] shape = new int[] {0};
 TensorBuffer tensorBuffer = TensorBuffer.createFixedSize(shape, DataType.UINT8);
 

The size of a fixed-size TensorBuffer cannot be changed once it is created.

Parameters
shape The shape of the TensorBuffer to be created.
dataType The dataType of the TensorBuffer to be created.
Throws
NullPointerException if shape is null.
IllegalArgumentException if shape has non-positive elements.

public static TensorBuffer createFrom (TensorBuffer buffer, DataType dataType)

Creates a TensorBuffer deep-copying data from another, with specified DataType.

Parameters
buffer the source TensorBuffer to copy from.
dataType the expected DataType of newly created TensorBuffer.
Throws
NullPointerException if buffer is null.

public ByteBuffer getBuffer ()

Returns the data buffer.

public abstract DataType getDataType ()

Returns the data type of this buffer.

public int getFlatSize ()

Gets the flatSize of the buffer.

Throws
IllegalStateException if the underlying data is corrupted

public abstract float[] getFloatArray ()

Returns a float array of the values stored in this buffer. If the buffer is of different types than float, the values will be converted into float. For example, values in TensorBufferUint8 will be converted from uint8 to float.

public abstract float getFloatValue (int absIndex)

Returns a float value at a given index. If the buffer is of different types than float, the value will be converted into float. For example, when reading a value from TensorBufferUint8, the value will be first read out as uint8, and then will be converted from uint8 to float.

 For example, a TensorBuffer with shape {2, 3} that represents the following array,
 [[0.0f, 1.0f, 2.0f], [3.0f, 4.0f, 5.0f]].

 The fourth element (whose value is 3.0f) in the TensorBuffer can be retrieved by:
 float v = tensorBuffer.getFloatValue(3);
 

Parameters
absIndex The absolute index of the value to be read.

public abstract int[] getIntArray ()

Returns an int array of the values stored in this buffer. If the buffer is of different type than int, the values will be converted into int, and loss of precision may apply. For example, getting an int array from a TensorBufferFloat with values {400.32f, 23.04f}, the output is {400, 23}.

public abstract int getIntValue (int absIndex)

Returns an int value at a given index. If the buffer is of different types than int, the value will be converted into int. For example, when reading a value from TensorBufferFloat, the value will be first read out as float, and then will be converted from float to int. Loss of precision may apply.

 For example, a TensorBuffer with shape {2, 3} that represents the following array,
 [[0.0f, 1.0f, 2.0f], [3.0f, 4.0f, 5.0f]].

 The fourth element (whose value is 3.0f) in the TensorBuffer can be retrieved by:
 int v = tensorBuffer.getIntValue(3);
 Note that v is converted from 3.0f to 3 as a result of type conversion.
 

Parameters
absIndex The absolute index of the value to be read.

public int[] getShape ()

Gets the current shape. (returning a copy here to avoid unexpected modification.)

Throws
IllegalStateException if the underlying data is corrupted

public abstract int getTypeSize ()

Returns the number of bytes of a single element in the array. For example, a float buffer will return 4, and a byte buffer will return 1.

public boolean isDynamic ()

Returns if the TensorBuffer is dynamic sized (could resize arbitrarily).

public abstract void loadArray (int[] src, int[] shape)

Loads an int array into this buffer with specific shape. If the buffer is of different types than int, the values will be converted into the buffer's type before being loaded into the buffer, and loss of precision may apply. For example, loading an int array with values {400, -23} into a TensorBufferUint8 , the values will be clamped to [0, 255] and then be casted to uint8 by {255, 0}.

Parameters
src The source array to be loaded.
shape Shape of the tensor that src represents.
Throws
NullPointerException if src is null.
NullPointerException if shape is null.
IllegalArgumentException if the size of the array to be loaded does not match the specified shape.

public abstract void loadArray (float[] src, int[] shape)

Loads a float array into this buffer with specific shape. If the buffer is of different types than float, the values will be converted into the buffer's type before being loaded into the buffer, and loss of precision may apply. For example, loading a float array into a TensorBufferUint8 with values {400.32f, -23.04f}, the values will be clamped to [0, 255] and then be casted to uint8 by {255, 0}.

Parameters
src The source array to be loaded.
shape Shape of the tensor that src represents.
Throws
NullPointerException if src is null.
NullPointerException if shape is null.
IllegalArgumentException if the size of the array to be loaded does not match the specified shape.

public void loadArray (float[] src)

Loads a float array into this buffer. If the buffer is of different types than float, the values will be converted into the buffer's type before being loaded into the buffer, and loss of precision may apply. For example, loading a float array into a TensorBufferUint8 with values {400.32f, -23.04f}, the values will be clamped to [0, 255] and then be casted to uint8 by {255, 0}.

Using this method assumes that the shape of src is the same as the shape of this TensorBuffer. Thus the size of buffer (src.length) should always match the flat size of this TensorBuffer, for both fixed-size and dynamic TensorBuffer. Use loadArray(float[], int[]) if src has a different shape.

Parameters
src The source array to be loaded.

public void loadArray (int[] src)

Loads an int array into this buffer. If the buffer is of different types than int, the values will be converted into the buffer's type before being loaded into the buffer, and loss of precision may apply. For example, loading an int array with values {400, -23} into a TensorBufferUint8 , the values will be clamped to [0, 255] and then be casted to uint8 by {255, 0}.

Using this method assumes that the shape of src is the same as the shape of this TensorBuffer. Thus the size of buffer (src.length) should always match the flat size of this TensorBuffer, for both fixed-size and dynamic TensorBuffer. Use loadArray(int[], int[]) if src has a different shape.

Parameters
src The source array to be loaded.

public void loadBuffer (ByteBuffer buffer)

Loads a byte buffer into this TensorBuffer. Buffer size must match the flat size of this TensorBuffer.

Using this method assumes that the shape of buffer is the same as the shape of this TensorBuffer. Thus the size of buffer (buffer.limit()) should always match the flat size of this TensorBuffer, for both fixed-size and dynamic TensorBuffer. Use loadBuffer(ByteBuffer, int[]) if buffer has a different shape.

Important: The loaded buffer is a reference. DO NOT MODIFY. We don't create a copy here for performance concern, but if modification is necessary, please make a copy.

For the best performance, always load a direct ByteBuffer or a ByteBuffer backed by an array.

If the buffer is read-only, we adopt a copy-on-write strategy for performance.

Parameters
buffer The byte buffer to load.

public void loadBuffer (ByteBuffer buffer, int[] shape)

Loads a byte buffer into this TensorBuffer with specific shape.

Important: The loaded buffer is a reference. DO NOT MODIFY. We don't create a copy here for performance concern, but if modification is necessary, please make a copy.

For the best performance, always load a direct ByteBuffer or a ByteBuffer backed by an array.

Parameters
buffer The byte buffer to load.
shape
Throws
NullPointerException if buffer is null.
IllegalArgumentException if the size of buffer and typeSize do not match or the size of buffer and flatSize do not match.