Mojo struct
ImageData
@register_passable(trivial)
struct ImageData[shape: DimList, type: DType, static_layout: Image2DLayout, origin: MutableOrigin]
Utility class that generalizes conv2d data and filter tensor with a given data layout.
Fields
- data (
NDBuffer[type, 4, origin, shape]): - dynamic_layout (
Image2DLayout):
Implemented traits
AnyType,
Copyable,
Movable,
UnknownDestructibility
Methods
__init__
__init__(data: NDBuffer[type, 4, origin, shape], layout: Image2DLayout) -> Self
Construct of an image data instance with dynamic layout param.
Args:
- data (
NDBuffer[type, 4, origin, shape]): A 4d buffer containing the actual data. - layout (
Image2DLayout): Data layout tag.
@implicit
__init__(data: NDBuffer[type, 4, origin, shape]) -> Self
__getitem__
__getitem__(self, n: Int, c: Int, h: Int, w: Int) -> SIMD[type, 1]
Reads the underlying data buffer based on the tensor index and under- lying data layout.
Args:
- n (
Int): Index on the batch dimension. - c (
Int): Index on the channel dimension. - h (
Int): Index on the height dimension. - w (
Int): Index on the width dimension.
Returns:
The value stored at the given index position.
__setitem__
__setitem__(self, n: Int, c: Int, h: Int, w: Int, value: SIMD[type, 1])
Writes the underlying data buffer based on the tensor index and under- lying data layout.
Args:
- n (
Int): Index on the batch dimension. - c (
Int): Index on the channel dimension. - h (
Int): Index on the height dimension. - w (
Int): Index on the width dimension. - value (
SIMD[type, 1]): The value to store at the given index position.
to_static_layout
to_static_layout[new_static_layout: Image2DLayout](self) -> ImageData[shape, type, new_static_layout, origin]
Conversion utility from a fully dynamic data structure, e.g. from c shim to one with compile-time known data layout.
Returns:
The image data with static data layout.
get_layout
get_layout(self) -> Image2DLayout
The getter function of the underlying data layout, resolving from either statically or dynamically provided information.
Returns:
The resolved data layout tag for this image instance.
get_flat_index
get_flat_index(self, n: Int, c: Int, h: Int, w: Int) -> Int
Converts the dimension index to the flat index of the underlying data based on the tensor layout.
Args:
- n (
Int): Index on the batch dimension. - c (
Int): Index on the channel dimension. - h (
Int): Index on the height dimension. - w (
Int): Index on the width dimension.
Returns:
An integer containing the index based on the underlying data layout.
get_tuple_index
get_tuple_index(self, idx: Int) -> IndexList[4]
Converts the flat index to the dimension index of the underlying data based on the tensor layout.
Args:
- idx (
Int): Flat index.
Returns:
A IndexList containing the index in NCHW order.
num_elements
num_elements(self) -> Int
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!