IMPORTANT: To view this page as Markdown, append `.md` to the URL (e.g. /get-started.md). For the complete documentation index, see llms.txt.
Skip to main content
For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).

Mojo trait

TensorOps

Extends TensorEngine with elementwise arithmetic.

A conforming type provides the same non-owning storage handle as TensorEngine, plus families of in-place (i*) and out-of-place elementwise operations. Binary and out-of-place unary operations take their operands as (storage, layout) tuples, so the layout describing a handle travels alongside it. Out-of-place operations are keyword-only with dst first.

Implemented traits​

AnyType, TensorEngine

Required methods​

iadd​

static def iadd[SelfLayoutType: TensorLayout, self_origin: MutOrigin, self_address_space: AddressSpace, OtherLayoutType: TensorLayout, other_mut: Bool, other_origin: Origin[mut=other_mut], other_address_space: AddressSpace, //, dtype: DType, OtherEngine: TensorEngine](storage: Tuple[Self.StorageType[True, origin_of(self_origin), dtype, self_origin, self_address_space], SelfLayoutType], other: Tuple[OtherEngine.StorageType[other_mut, origin_of(other_origin), dtype, other_origin, other_address_space], OtherLayoutType])

Adds other into storage elementwise, in place.

Parameters:

  • ​SelfLayoutType (TensorLayout): The layout type of the destination storage.
  • ​self_origin (MutOrigin): The origin of the destination storage.
  • ​self_address_space (AddressSpace): The address space of the destination storage.
  • ​OtherLayoutType (TensorLayout): The layout type of the right-hand storage operand.
  • ​other_mut (Bool): The mutability of the right-hand storage operand.
  • ​other_origin (Origin[mut=other_mut]): The origin of the right-hand storage operand.
  • ​other_address_space (AddressSpace): The address space of the right-hand storage operand.
  • ​dtype (DType): The element data type of both storages.
  • ​OtherEngine (TensorEngine): The engine of the right-hand operand. May differ from Self as long as the two engines have the same logical element size.

Args:

imul​

static def imul[SelfLayoutType: TensorLayout, self_origin: MutOrigin, self_address_space: AddressSpace, OtherLayoutType: TensorLayout, other_mut: Bool, other_origin: Origin[mut=other_mut], other_address_space: AddressSpace, //, dtype: DType, OtherEngine: TensorEngine](storage: Tuple[Self.StorageType[True, origin_of(self_origin), dtype, self_origin, self_address_space], SelfLayoutType], other: Tuple[OtherEngine.StorageType[other_mut, origin_of(other_origin), dtype, other_origin, other_address_space], OtherLayoutType])

Multiplies storage by other elementwise, in place.

Parameters:

  • ​SelfLayoutType (TensorLayout): The layout type of the destination storage.
  • ​self_origin (MutOrigin): The origin of the destination storage.
  • ​self_address_space (AddressSpace): The address space of the destination storage.
  • ​OtherLayoutType (TensorLayout): The layout type of the right-hand storage operand.
  • ​other_mut (Bool): The mutability of the right-hand storage operand.
  • ​other_origin (Origin[mut=other_mut]): The origin of the right-hand storage operand.
  • ​other_address_space (AddressSpace): The address space of the right-hand storage operand.
  • ​dtype (DType): The element data type of both storages.
  • ​OtherEngine (TensorEngine): The engine of the right-hand operand. May differ from Self as long as the two engines have the same logical element size.

Args:

isub​

static def isub[SelfLayoutType: TensorLayout, self_origin: MutOrigin, self_address_space: AddressSpace, OtherLayoutType: TensorLayout, other_mut: Bool, other_origin: Origin[mut=other_mut], other_address_space: AddressSpace, //, dtype: DType, OtherEngine: TensorEngine](storage: Tuple[Self.StorageType[True, origin_of(self_origin), dtype, self_origin, self_address_space], SelfLayoutType], other: Tuple[OtherEngine.StorageType[other_mut, origin_of(other_origin), dtype, other_origin, other_address_space], OtherLayoutType])

Subtracts other from storage elementwise, in place.

Parameters:

  • ​SelfLayoutType (TensorLayout): The layout type of the destination storage.
  • ​self_origin (MutOrigin): The origin of the destination storage.
  • ​self_address_space (AddressSpace): The address space of the destination storage.
  • ​OtherLayoutType (TensorLayout): The layout type of the right-hand storage operand.
  • ​other_mut (Bool): The mutability of the right-hand storage operand.
  • ​other_origin (Origin[mut=other_mut]): The origin of the right-hand storage operand.
  • ​other_address_space (AddressSpace): The address space of the right-hand storage operand.
  • ​dtype (DType): The element data type of both storages.
  • ​OtherEngine (TensorEngine): The engine of the right-hand operand. May differ from Self as long as the two engines have the same logical element size.

Args:

ifloordiv​

static def ifloordiv[SelfLayoutType: TensorLayout, self_origin: MutOrigin, self_address_space: AddressSpace, OtherLayoutType: TensorLayout, other_mut: Bool, other_origin: Origin[mut=other_mut], other_address_space: AddressSpace, //, dtype: DType, OtherEngine: TensorEngine](storage: Tuple[Self.StorageType[True, origin_of(self_origin), dtype, self_origin, self_address_space], SelfLayoutType], other: Tuple[OtherEngine.StorageType[other_mut, origin_of(other_origin), dtype, other_origin, other_address_space], OtherLayoutType])

Floor-divides storage by other elementwise, in place.

Parameters:

  • ​SelfLayoutType (TensorLayout): The layout type of the destination storage.
  • ​self_origin (MutOrigin): The origin of the destination storage.
  • ​self_address_space (AddressSpace): The address space of the destination storage.
  • ​OtherLayoutType (TensorLayout): The layout type of the right-hand storage operand.
  • ​other_mut (Bool): The mutability of the right-hand storage operand.
  • ​other_origin (Origin[mut=other_mut]): The origin of the right-hand storage operand.
  • ​other_address_space (AddressSpace): The address space of the right-hand storage operand.
  • ​dtype (DType): The element data type of both storages.
  • ​OtherEngine (TensorEngine): The engine of the right-hand operand. May differ from Self as long as the two engines have the same logical element size.

Args:

itruediv​

static def itruediv[SelfLayoutType: TensorLayout, self_origin: MutOrigin, self_address_space: AddressSpace, OtherLayoutType: TensorLayout, other_mut: Bool, other_origin: Origin[mut=other_mut], other_address_space: AddressSpace, //, dtype: DType, OtherEngine: TensorEngine](storage: Tuple[Self.StorageType[True, origin_of(self_origin), dtype, self_origin, self_address_space], SelfLayoutType], other: Tuple[OtherEngine.StorageType[other_mut, origin_of(other_origin), dtype, other_origin, other_address_space], OtherLayoutType])

True-divides storage by other elementwise, in place.

Parameters:

  • ​SelfLayoutType (TensorLayout): The layout type of the destination storage.
  • ​self_origin (MutOrigin): The origin of the destination storage.
  • ​self_address_space (AddressSpace): The address space of the destination storage.
  • ​OtherLayoutType (TensorLayout): The layout type of the right-hand storage operand.
  • ​other_mut (Bool): The mutability of the right-hand storage operand.
  • ​other_origin (Origin[mut=other_mut]): The origin of the right-hand storage operand.
  • ​other_address_space (AddressSpace): The address space of the right-hand storage operand.
  • ​dtype (DType): The element data type of both storages.
  • ​OtherEngine (TensorEngine): The engine of the right-hand operand. May differ from Self as long as the two engines have the same logical element size.

Args:

imin​

static def imin[SelfLayoutType: TensorLayout, self_origin: MutOrigin, self_address_space: AddressSpace, OtherLayoutType: TensorLayout, other_mut: Bool, other_origin: Origin[mut=other_mut], other_address_space: AddressSpace, //, dtype: DType, OtherEngine: TensorEngine](storage: Tuple[Self.StorageType[True, origin_of(self_origin), dtype, self_origin, self_address_space], SelfLayoutType], other: Tuple[OtherEngine.StorageType[other_mut, origin_of(other_origin), dtype, other_origin, other_address_space], OtherLayoutType])

Takes the elementwise minimum of storage and other, in place.

Parameters:

  • ​SelfLayoutType (TensorLayout): The layout type of the destination storage.
  • ​self_origin (MutOrigin): The origin of the destination storage.
  • ​self_address_space (AddressSpace): The address space of the destination storage.
  • ​OtherLayoutType (TensorLayout): The layout type of the right-hand storage operand.
  • ​other_mut (Bool): The mutability of the right-hand storage operand.
  • ​other_origin (Origin[mut=other_mut]): The origin of the right-hand storage operand.
  • ​other_address_space (AddressSpace): The address space of the right-hand storage operand.
  • ​dtype (DType): The element data type of both storages.
  • ​OtherEngine (TensorEngine): The engine of the right-hand operand. May differ from Self as long as the two engines have the same logical element size.

Args:

imax​

static def imax[SelfLayoutType: TensorLayout, self_origin: MutOrigin, self_address_space: AddressSpace, OtherLayoutType: TensorLayout, other_mut: Bool, other_origin: Origin[mut=other_mut], other_address_space: AddressSpace, //, dtype: DType, OtherEngine: TensorEngine](storage: Tuple[Self.StorageType[True, origin_of(self_origin), dtype, self_origin, self_address_space], SelfLayoutType], other: Tuple[OtherEngine.StorageType[other_mut, origin_of(other_origin), dtype, other_origin, other_address_space], OtherLayoutType])

Takes the elementwise maximum of storage and other, in place.

Parameters:

  • ​SelfLayoutType (TensorLayout): The layout type of the destination storage.
  • ​self_origin (MutOrigin): The origin of the destination storage.
  • ​self_address_space (AddressSpace): The address space of the destination storage.
  • ​OtherLayoutType (TensorLayout): The layout type of the right-hand storage operand.
  • ​other_mut (Bool): The mutability of the right-hand storage operand.
  • ​other_origin (Origin[mut=other_mut]): The origin of the right-hand storage operand.
  • ​other_address_space (AddressSpace): The address space of the right-hand storage operand.
  • ​dtype (DType): The element data type of both storages.
  • ​OtherEngine (TensorEngine): The engine of the right-hand operand. May differ from Self as long as the two engines have the same logical element size.

Args:

iabs​

static def iabs[dtype: DType, //](storage: Self.StorageType[True, origin_of(origin), dtype, origin, address_space], layout: T)

Takes the elementwise absolute value of storage, in place.

For unsigned dtypes this is the identity.

Parameters:

  • ​dtype (DType): The element data type of the storage.

Args:

  • ​storage (_Self.StorageType[True, origin_of(origin), dtype, origin, address_space]): The storage to modify in place.
  • ​layout (T): The layout describing the storage's elements.

irecip​

static def irecip[dtype: DType, //](storage: Self.StorageType[True, origin_of(origin), dtype, origin, address_space], layout: T)

Replaces each element of storage with its reciprocal, in place.

Elements equal to zero produce infinity, following IEEE 754 division semantics.

Parameters:

  • ​dtype (DType): The element data type of the storage. Must be a floating-point type.

Args:

  • ​storage (_Self.StorageType[True, origin_of(origin), dtype, origin, address_space]): The storage to modify in place.
  • ​layout (T): The layout describing the storage's elements.

iexp​

static def iexp[dtype: DType, scale_dtype: DType = dtype, //, scale: Scalar[scale_dtype]](storage: Self.StorageType[True, origin_of(origin), dtype, origin, address_space], layout: T)

Replaces each element x of storage with exp(scale * x), in place.

The scale factor is applied before exponentiation so that scaled exponentials (for example softmax logit scaling) fuse into a single pass over the elements. Pass a scale of 1 for a plain exponential.

Parameters:

  • ​dtype (DType): The element data type of the storage. Must be a floating-point type.
  • ​scale_dtype (DType): The data type of the scale factor. Defaults to dtype; the scale is cast to dtype before the multiplication.
  • ​scale (Scalar[scale_dtype]): The compile-time factor each element is multiplied by before exponentiation.

Args:

  • ​storage (_Self.StorageType[True, origin_of(origin), dtype, origin, address_space]): The storage to modify in place.
  • ​layout (T): The layout describing the storage's elements.

add​

static def add[DstLayoutType: TensorLayout, dst_origin: MutOrigin, dst_address_space: AddressSpace, LhsLayoutType: TensorLayout, lhs_mut: Bool, lhs_origin: Origin[mut=lhs_mut], lhs_address_space: AddressSpace, RhsLayoutType: TensorLayout, rhs_mut: Bool, rhs_origin: Origin[mut=rhs_mut], rhs_address_space: AddressSpace, //, dtype: DType, LhsEngine: TensorEngine, RhsEngine: TensorEngine](*, dst: Tuple[Self.StorageType[True, origin_of(dst_origin), dtype, dst_origin, dst_address_space], DstLayoutType], lhs: Tuple[LhsEngine.StorageType[lhs_mut, origin_of(lhs_origin), dtype, lhs_origin, lhs_address_space], LhsLayoutType], rhs: Tuple[RhsEngine.StorageType[rhs_mut, origin_of(rhs_origin), dtype, rhs_origin, rhs_address_space], RhsLayoutType])

Adds lhs and rhs elementwise, writing into dst.

Parameters:

  • ​DstLayoutType (TensorLayout): The layout type of the destination storage.
  • ​dst_origin (MutOrigin): The origin of the destination storage.
  • ​dst_address_space (AddressSpace): The address space of the destination storage.
  • ​LhsLayoutType (TensorLayout): The layout type of the left-hand storage operand.
  • ​lhs_mut (Bool): The mutability of the left-hand storage operand.
  • ​lhs_origin (Origin[mut=lhs_mut]): The origin of the left-hand storage operand.
  • ​lhs_address_space (AddressSpace): The address space of the left-hand storage operand.
  • ​RhsLayoutType (TensorLayout): The layout type of the right-hand storage operand.
  • ​rhs_mut (Bool): The mutability of the right-hand storage operand.
  • ​rhs_origin (Origin[mut=rhs_mut]): The origin of the right-hand storage operand.
  • ​rhs_address_space (AddressSpace): The address space of the right-hand storage operand.
  • ​dtype (DType): The element data type of all three storages.
  • ​LhsEngine (TensorEngine): The engine of the left-hand operand. May differ from Self as long as the engines have the same logical element size.
  • ​RhsEngine (TensorEngine): The engine of the right-hand operand. May differ from Self as long as the engines have the same logical element size.

Args:

mul​

static def mul[DstLayoutType: TensorLayout, dst_origin: MutOrigin, dst_address_space: AddressSpace, LhsLayoutType: TensorLayout, lhs_mut: Bool, lhs_origin: Origin[mut=lhs_mut], lhs_address_space: AddressSpace, RhsLayoutType: TensorLayout, rhs_mut: Bool, rhs_origin: Origin[mut=rhs_mut], rhs_address_space: AddressSpace, //, dtype: DType, LhsEngine: TensorEngine, RhsEngine: TensorEngine](*, dst: Tuple[Self.StorageType[True, origin_of(dst_origin), dtype, dst_origin, dst_address_space], DstLayoutType], lhs: Tuple[LhsEngine.StorageType[lhs_mut, origin_of(lhs_origin), dtype, lhs_origin, lhs_address_space], LhsLayoutType], rhs: Tuple[RhsEngine.StorageType[rhs_mut, origin_of(rhs_origin), dtype, rhs_origin, rhs_address_space], RhsLayoutType])

Multiplies lhs by rhs elementwise, writing into dst.

Parameters:

  • ​DstLayoutType (TensorLayout): The layout type of the destination storage.
  • ​dst_origin (MutOrigin): The origin of the destination storage.
  • ​dst_address_space (AddressSpace): The address space of the destination storage.
  • ​LhsLayoutType (TensorLayout): The layout type of the left-hand storage operand.
  • ​lhs_mut (Bool): The mutability of the left-hand storage operand.
  • ​lhs_origin (Origin[mut=lhs_mut]): The origin of the left-hand storage operand.
  • ​lhs_address_space (AddressSpace): The address space of the left-hand storage operand.
  • ​RhsLayoutType (TensorLayout): The layout type of the right-hand storage operand.
  • ​rhs_mut (Bool): The mutability of the right-hand storage operand.
  • ​rhs_origin (Origin[mut=rhs_mut]): The origin of the right-hand storage operand.
  • ​rhs_address_space (AddressSpace): The address space of the right-hand storage operand.
  • ​dtype (DType): The element data type of all three storages.
  • ​LhsEngine (TensorEngine): The engine of the left-hand operand. May differ from Self as long as the engines have the same logical element size.
  • ​RhsEngine (TensorEngine): The engine of the right-hand operand. May differ from Self as long as the engines have the same logical element size.

Args:

sub​

static def sub[DstLayoutType: TensorLayout, dst_origin: MutOrigin, dst_address_space: AddressSpace, LhsLayoutType: TensorLayout, lhs_mut: Bool, lhs_origin: Origin[mut=lhs_mut], lhs_address_space: AddressSpace, RhsLayoutType: TensorLayout, rhs_mut: Bool, rhs_origin: Origin[mut=rhs_mut], rhs_address_space: AddressSpace, //, dtype: DType, LhsEngine: TensorEngine, RhsEngine: TensorEngine](*, dst: Tuple[Self.StorageType[True, origin_of(dst_origin), dtype, dst_origin, dst_address_space], DstLayoutType], lhs: Tuple[LhsEngine.StorageType[lhs_mut, origin_of(lhs_origin), dtype, lhs_origin, lhs_address_space], LhsLayoutType], rhs: Tuple[RhsEngine.StorageType[rhs_mut, origin_of(rhs_origin), dtype, rhs_origin, rhs_address_space], RhsLayoutType])

Subtracts rhs from lhs elementwise, writing into dst.

Parameters:

  • ​DstLayoutType (TensorLayout): The layout type of the destination storage.
  • ​dst_origin (MutOrigin): The origin of the destination storage.
  • ​dst_address_space (AddressSpace): The address space of the destination storage.
  • ​LhsLayoutType (TensorLayout): The layout type of the left-hand storage operand.
  • ​lhs_mut (Bool): The mutability of the left-hand storage operand.
  • ​lhs_origin (Origin[mut=lhs_mut]): The origin of the left-hand storage operand.
  • ​lhs_address_space (AddressSpace): The address space of the left-hand storage operand.
  • ​RhsLayoutType (TensorLayout): The layout type of the right-hand storage operand.
  • ​rhs_mut (Bool): The mutability of the right-hand storage operand.
  • ​rhs_origin (Origin[mut=rhs_mut]): The origin of the right-hand storage operand.
  • ​rhs_address_space (AddressSpace): The address space of the right-hand storage operand.
  • ​dtype (DType): The element data type of all three storages.
  • ​LhsEngine (TensorEngine): The engine of the left-hand operand. May differ from Self as long as the engines have the same logical element size.
  • ​RhsEngine (TensorEngine): The engine of the right-hand operand. May differ from Self as long as the engines have the same logical element size.

Args:

floordiv​

static def floordiv[DstLayoutType: TensorLayout, dst_origin: MutOrigin, dst_address_space: AddressSpace, LhsLayoutType: TensorLayout, lhs_mut: Bool, lhs_origin: Origin[mut=lhs_mut], lhs_address_space: AddressSpace, RhsLayoutType: TensorLayout, rhs_mut: Bool, rhs_origin: Origin[mut=rhs_mut], rhs_address_space: AddressSpace, //, dtype: DType, LhsEngine: TensorEngine, RhsEngine: TensorEngine](*, dst: Tuple[Self.StorageType[True, origin_of(dst_origin), dtype, dst_origin, dst_address_space], DstLayoutType], lhs: Tuple[LhsEngine.StorageType[lhs_mut, origin_of(lhs_origin), dtype, lhs_origin, lhs_address_space], LhsLayoutType], rhs: Tuple[RhsEngine.StorageType[rhs_mut, origin_of(rhs_origin), dtype, rhs_origin, rhs_address_space], RhsLayoutType])

Floor-divides lhs by rhs elementwise, writing into dst.

Parameters:

  • ​DstLayoutType (TensorLayout): The layout type of the destination storage.
  • ​dst_origin (MutOrigin): The origin of the destination storage.
  • ​dst_address_space (AddressSpace): The address space of the destination storage.
  • ​LhsLayoutType (TensorLayout): The layout type of the left-hand storage operand.
  • ​lhs_mut (Bool): The mutability of the left-hand storage operand.
  • ​lhs_origin (Origin[mut=lhs_mut]): The origin of the left-hand storage operand.
  • ​lhs_address_space (AddressSpace): The address space of the left-hand storage operand.
  • ​RhsLayoutType (TensorLayout): The layout type of the right-hand storage operand.
  • ​rhs_mut (Bool): The mutability of the right-hand storage operand.
  • ​rhs_origin (Origin[mut=rhs_mut]): The origin of the right-hand storage operand.
  • ​rhs_address_space (AddressSpace): The address space of the right-hand storage operand.
  • ​dtype (DType): The element data type of all three storages.
  • ​LhsEngine (TensorEngine): The engine of the left-hand operand. May differ from Self as long as the engines have the same logical element size.
  • ​RhsEngine (TensorEngine): The engine of the right-hand operand. May differ from Self as long as the engines have the same logical element size.

Args:

truediv​

static def truediv[DstLayoutType: TensorLayout, dst_origin: MutOrigin, dst_address_space: AddressSpace, LhsLayoutType: TensorLayout, lhs_mut: Bool, lhs_origin: Origin[mut=lhs_mut], lhs_address_space: AddressSpace, RhsLayoutType: TensorLayout, rhs_mut: Bool, rhs_origin: Origin[mut=rhs_mut], rhs_address_space: AddressSpace, //, dtype: DType, LhsEngine: TensorEngine, RhsEngine: TensorEngine](*, dst: Tuple[Self.StorageType[True, origin_of(dst_origin), dtype, dst_origin, dst_address_space], DstLayoutType], lhs: Tuple[LhsEngine.StorageType[lhs_mut, origin_of(lhs_origin), dtype, lhs_origin, lhs_address_space], LhsLayoutType], rhs: Tuple[RhsEngine.StorageType[rhs_mut, origin_of(rhs_origin), dtype, rhs_origin, rhs_address_space], RhsLayoutType])

True-divides lhs by rhs elementwise, writing into dst.

Parameters:

  • ​DstLayoutType (TensorLayout): The layout type of the destination storage.
  • ​dst_origin (MutOrigin): The origin of the destination storage.
  • ​dst_address_space (AddressSpace): The address space of the destination storage.
  • ​LhsLayoutType (TensorLayout): The layout type of the left-hand storage operand.
  • ​lhs_mut (Bool): The mutability of the left-hand storage operand.
  • ​lhs_origin (Origin[mut=lhs_mut]): The origin of the left-hand storage operand.
  • ​lhs_address_space (AddressSpace): The address space of the left-hand storage operand.
  • ​RhsLayoutType (TensorLayout): The layout type of the right-hand storage operand.
  • ​rhs_mut (Bool): The mutability of the right-hand storage operand.
  • ​rhs_origin (Origin[mut=rhs_mut]): The origin of the right-hand storage operand.
  • ​rhs_address_space (AddressSpace): The address space of the right-hand storage operand.
  • ​dtype (DType): The element data type of all three storages.
  • ​LhsEngine (TensorEngine): The engine of the left-hand operand. May differ from Self as long as the engines have the same logical element size.
  • ​RhsEngine (TensorEngine): The engine of the right-hand operand. May differ from Self as long as the engines have the same logical element size.

Args:

min​

static def min[DstLayoutType: TensorLayout, dst_origin: MutOrigin, dst_address_space: AddressSpace, LhsLayoutType: TensorLayout, lhs_mut: Bool, lhs_origin: Origin[mut=lhs_mut], lhs_address_space: AddressSpace, RhsLayoutType: TensorLayout, rhs_mut: Bool, rhs_origin: Origin[mut=rhs_mut], rhs_address_space: AddressSpace, //, dtype: DType, LhsEngine: TensorEngine, RhsEngine: TensorEngine](*, dst: Tuple[Self.StorageType[True, origin_of(dst_origin), dtype, dst_origin, dst_address_space], DstLayoutType], lhs: Tuple[LhsEngine.StorageType[lhs_mut, origin_of(lhs_origin), dtype, lhs_origin, lhs_address_space], LhsLayoutType], rhs: Tuple[RhsEngine.StorageType[rhs_mut, origin_of(rhs_origin), dtype, rhs_origin, rhs_address_space], RhsLayoutType])

Takes the elementwise minimum of lhs and rhs, writing into dst.

Parameters:

  • ​DstLayoutType (TensorLayout): The layout type of the destination storage.
  • ​dst_origin (MutOrigin): The origin of the destination storage.
  • ​dst_address_space (AddressSpace): The address space of the destination storage.
  • ​LhsLayoutType (TensorLayout): The layout type of the left-hand storage operand.
  • ​lhs_mut (Bool): The mutability of the left-hand storage operand.
  • ​lhs_origin (Origin[mut=lhs_mut]): The origin of the left-hand storage operand.
  • ​lhs_address_space (AddressSpace): The address space of the left-hand storage operand.
  • ​RhsLayoutType (TensorLayout): The layout type of the right-hand storage operand.
  • ​rhs_mut (Bool): The mutability of the right-hand storage operand.
  • ​rhs_origin (Origin[mut=rhs_mut]): The origin of the right-hand storage operand.
  • ​rhs_address_space (AddressSpace): The address space of the right-hand storage operand.
  • ​dtype (DType): The element data type of all three storages.
  • ​LhsEngine (TensorEngine): The engine of the left-hand operand. May differ from Self as long as the engines have the same logical element size.
  • ​RhsEngine (TensorEngine): The engine of the right-hand operand. May differ from Self as long as the engines have the same logical element size.

Args:

max​

static def max[DstLayoutType: TensorLayout, dst_origin: MutOrigin, dst_address_space: AddressSpace, LhsLayoutType: TensorLayout, lhs_mut: Bool, lhs_origin: Origin[mut=lhs_mut], lhs_address_space: AddressSpace, RhsLayoutType: TensorLayout, rhs_mut: Bool, rhs_origin: Origin[mut=rhs_mut], rhs_address_space: AddressSpace, //, dtype: DType, LhsEngine: TensorEngine, RhsEngine: TensorEngine](*, dst: Tuple[Self.StorageType[True, origin_of(dst_origin), dtype, dst_origin, dst_address_space], DstLayoutType], lhs: Tuple[LhsEngine.StorageType[lhs_mut, origin_of(lhs_origin), dtype, lhs_origin, lhs_address_space], LhsLayoutType], rhs: Tuple[RhsEngine.StorageType[rhs_mut, origin_of(rhs_origin), dtype, rhs_origin, rhs_address_space], RhsLayoutType])

Takes the elementwise maximum of lhs and rhs, writing into dst.

Parameters:

  • ​DstLayoutType (TensorLayout): The layout type of the destination storage.
  • ​dst_origin (MutOrigin): The origin of the destination storage.
  • ​dst_address_space (AddressSpace): The address space of the destination storage.
  • ​LhsLayoutType (TensorLayout): The layout type of the left-hand storage operand.
  • ​lhs_mut (Bool): The mutability of the left-hand storage operand.
  • ​lhs_origin (Origin[mut=lhs_mut]): The origin of the left-hand storage operand.
  • ​lhs_address_space (AddressSpace): The address space of the left-hand storage operand.
  • ​RhsLayoutType (TensorLayout): The layout type of the right-hand storage operand.
  • ​rhs_mut (Bool): The mutability of the right-hand storage operand.
  • ​rhs_origin (Origin[mut=rhs_mut]): The origin of the right-hand storage operand.
  • ​rhs_address_space (AddressSpace): The address space of the right-hand storage operand.
  • ​dtype (DType): The element data type of all three storages.
  • ​LhsEngine (TensorEngine): The engine of the left-hand operand. May differ from Self as long as the engines have the same logical element size.
  • ​RhsEngine (TensorEngine): The engine of the right-hand operand. May differ from Self as long as the engines have the same logical element size.

Args:

abs​

static def abs[DstLayoutType: TensorLayout, dst_origin: MutOrigin, dst_address_space: AddressSpace, SrcLayoutType: TensorLayout, src_mut: Bool, src_origin: Origin[mut=src_mut], src_address_space: AddressSpace, //, dtype: DType, SrcEngine: TensorEngine](*, dst: Tuple[Self.StorageType[True, origin_of(dst_origin), dtype, dst_origin, dst_address_space], DstLayoutType], src: Tuple[SrcEngine.StorageType[src_mut, origin_of(src_origin), dtype, src_origin, src_address_space], SrcLayoutType])

Takes the elementwise absolute value of src, writing into dst.

For unsigned dtypes this is the identity.

Parameters:

  • ​DstLayoutType (TensorLayout): The layout type of the destination storage.
  • ​dst_origin (MutOrigin): The origin of the destination storage.
  • ​dst_address_space (AddressSpace): The address space of the destination storage.
  • ​SrcLayoutType (TensorLayout): The layout type of the source storage.
  • ​src_mut (Bool): The mutability of the source storage.
  • ​src_origin (Origin[mut=src_mut]): The origin of the source storage.
  • ​src_address_space (AddressSpace): The address space of the source storage.
  • ​dtype (DType): The element data type of both storages.
  • ​SrcEngine (TensorEngine): The engine of the source. May differ from Self as long as the two engines have the same logical element size.

Args:

recip​

static def recip[DstLayoutType: TensorLayout, dst_origin: MutOrigin, dst_address_space: AddressSpace, SrcLayoutType: TensorLayout, src_mut: Bool, src_origin: Origin[mut=src_mut], src_address_space: AddressSpace, //, dtype: DType, SrcEngine: TensorEngine](*, dst: Tuple[Self.StorageType[True, origin_of(dst_origin), dtype, dst_origin, dst_address_space], DstLayoutType], src: Tuple[SrcEngine.StorageType[src_mut, origin_of(src_origin), dtype, src_origin, src_address_space], SrcLayoutType])

Writes the reciprocal of each element of src into dst.

Elements equal to zero produce infinity, following IEEE 754 division semantics.

Parameters:

  • ​DstLayoutType (TensorLayout): The layout type of the destination storage.
  • ​dst_origin (MutOrigin): The origin of the destination storage.
  • ​dst_address_space (AddressSpace): The address space of the destination storage.
  • ​SrcLayoutType (TensorLayout): The layout type of the source storage.
  • ​src_mut (Bool): The mutability of the source storage.
  • ​src_origin (Origin[mut=src_mut]): The origin of the source storage.
  • ​src_address_space (AddressSpace): The address space of the source storage.
  • ​dtype (DType): The element data type of both storages.
  • ​SrcEngine (TensorEngine): The engine of the source. May differ from Self as long as the two engines have the same logical element size.

Args:

exp​

static def exp[DstLayoutType: TensorLayout, dst_origin: MutOrigin, dst_address_space: AddressSpace, SrcLayoutType: TensorLayout, src_mut: Bool, src_origin: Origin[mut=src_mut], src_address_space: AddressSpace, dtype: DType, SrcEngine: TensorEngine, scale_dtype: DType = dtype, //, scale: Scalar[scale_dtype] = 1](*, dst: Tuple[Self.StorageType[True, origin_of(dst_origin), dtype, dst_origin, dst_address_space], DstLayoutType], src: Tuple[SrcEngine.StorageType[src_mut, origin_of(src_origin), dtype, src_origin, src_address_space], SrcLayoutType])

Writes exp(scale * x) for each element x of src into dst.

The scale factor is applied before exponentiation so that scaled exponentials (for example softmax logit scaling) fuse into a single pass over the elements. Pass a scale of 1 for a plain exponential.

Parameters:

  • ​DstLayoutType (TensorLayout): The layout type of the destination storage.
  • ​dst_origin (MutOrigin): The origin of the destination storage.
  • ​dst_address_space (AddressSpace): The address space of the destination storage.
  • ​SrcLayoutType (TensorLayout): The layout type of the source storage.
  • ​src_mut (Bool): The mutability of the source storage.
  • ​src_origin (Origin[mut=src_mut]): The origin of the source storage.
  • ​src_address_space (AddressSpace): The address space of the source storage.
  • ​dtype (DType): The element data type of both storages. Must be a floating-point type.
  • ​SrcEngine (TensorEngine): The engine of the source. May differ from Self as long as the two engines have the same logical element size.
  • ​scale_dtype (DType): The data type of the scale factor. Defaults to dtype; the scale is cast to dtype before the multiplication.
  • ​scale (Scalar[scale_dtype]): The compile-time factor each element is multiplied by before exponentiation.

Args:

View source

View source

Was this page helpful?

Models

Cookbook

Puzzles

Blog

Community