Skip to main content
Log in

Mojo struct

MaterializedMask

@register_passable(trivial) struct MaterializedMask[type_: DType, rank_: Int, shape_: DimList]

Mask that's backed by a materialized tensor.

Aliases

  • apply_log2e_after_mask = True:
  • mask_out_of_bound = True:
  • mask_safe_out_of_bounds = False:
  • MaskType = NDBuffer[type_, rank_, MutableAnyOrigin, shape_]:

Fields

  • mask_tensor (NDBuffer[type_, rank_, MutableAnyOrigin, shape_]):
  • start_pos (OptionalReg[NDBuffer[uint32, 1, MutableAnyOrigin]]):
  • is_multiple_of_2 (Bool):

Implemented traits

AnyType, Copyable, ExplicitlyCopyable, MHAMask, Movable, UnknownDestructibility

Methods

__init__

__init__(mask_tensor: NDBuffer[type_, rank_, MutableAnyOrigin, shape_], start_pos: OptionalReg[NDBuffer[uint32, 1, MutableAnyOrigin]] = OptionalReg[NDBuffer[uint32, 1, MutableAnyOrigin]]({:i1 0, 1})) -> Self

get_start_pos

get_start_pos(self, batch_idx: Int) -> Int

mask

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

status

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

Was this page helpful?