Skip to main content

Mojo struct

MaterializedMask

@register_passable(trivial) struct MaterializedMask[dtype_: DType, layout_: Layout]

Mask that's backed by a materialized tensor.

Fields

  • mask_tensor (LayoutTensor[dtype_, layout_, MutableAnyOrigin]):
  • start_pos (OptionalReg[LayoutTensor[DType.uint32, Layout.row_major(-1), MutableAnyOrigin]]):
  • is_multiple_of_2 (Bool):

Implemented traits

AnyType, Copyable, DevicePassable, ImplicitlyCopyable, MHAMask, Movable, UnknownDestructibility

Aliases

__copyinit__is_trivial

alias __copyinit__is_trivial = True

__del__is_trivial

alias __del__is_trivial = True

__moveinit__is_trivial

alias __moveinit__is_trivial = True

apply_log2e_after_mask

alias apply_log2e_after_mask = True

check_mask_during_decoding

alias check_mask_during_decoding = True

device_type

alias device_type = MaterializedMask[dtype_, layout_]

mask_out_of_bound

alias mask_out_of_bound = True

mask_safe_out_of_bounds

alias mask_safe_out_of_bounds = False

MaskType

alias MaskType = LayoutTensor[dtype_, layout_, MutableAnyOrigin]

Methods

__init__

__init__(mask_tensor: LayoutTensor[dtype_, layout_, MutableAnyOrigin], start_pos: OptionalReg[LayoutTensor[DType.uint32, Layout.row_major(-1), MutableAnyOrigin]] = None) -> Self

get_type_name

static get_type_name() -> String

Returns:

String

get_device_type_name

static get_device_type_name() -> String

Returns:

String

get_start_pos

get_start_pos(self, batch_idx: Int) -> Int

Returns:

Int

mask

mask[dtype: DType, width: Int, //, *, element_type: DType = DType.uint32](self, coord: IndexList[4, element_type=element_type], score_vec: SIMD[dtype, width]) -> SIMD[dtype, width]

Returns:

SIMD

status

status[*, element_type: DType = DType.uint32](self, tile_offset: IndexList[2, element_type=element_type], tile_size: IndexList[2, element_type=element_type]) -> TileMaskStatus

Returns:

TileMaskStatus

Was this page helpful?