Skip to main content
Log in

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, 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 staticall or dynamicall 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) -> Index[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?