tenzr::plax¶

Classes¶

Functions¶

Functions Documentation¶

function Add¶

TENZR_APIArray Add(
    const Array & lhs,
    const Array & rhs
)

Adds two tensors of identical dtype elementwise.

Parameters:

• lhs Left operand.
• rhs Right operand.

Return: lhs + rhs.

Primitive Linear Algebra for XLA (plax), analogous to JAX's jax.lax. Provides a low-level library of primitive operations that serve as the foundation for higher-level APIs such as tenzr::numpy. Key principles: - Zero implicit type promotion: binary operators require strictly matching dtypes. - Zero implicit broadcasting: dimensions must match unless explicitly broadcast. - 1:1 mapping to OpenXLA / StableHLO compiler primitives.

Strict elementwise addition; operands must share dtype and shape.

function Sub¶

TENZR_APIArray Sub(
    const Array & lhs,
    const Array & rhs
)

Subtracts rhs from lhs elementwise.

Parameters:

• lhs Left operand.
• rhs Right operand.

Return: lhs - rhs.

Strict elementwise subtraction; operands must share dtype and shape.

function Mul¶

TENZR_APIArray Mul(
    const Array & lhs,
    const Array & rhs
)

Multiplies two tensors elementwise.

Parameters:

• lhs Left operand.
• rhs Right operand.

Return: lhs * rhs.

Strict elementwise multiplication; operands must share dtype and shape.

function Div¶

TENZR_APIArray Div(
    const Array & lhs,
    const Array & rhs
)

Divides lhs by rhs elementwise.

Parameters:

• lhs Dividend.
• rhs Divisor.

Return: lhs / rhs.

Strict elementwise division; operands must share dtype and shape.

function Pow¶

TENZR_APIArray Pow(
    const Array & lhs,
    const Array & rhs
)

Raises lhs to the power rhs elementwise.

Parameters:

• lhs Base.
• rhs Exponent.

Return: lhs ** rhs.

Strict elementwise exponentiation; operands must share dtype and shape.

function Maximum¶

TENZR_APIArray Maximum(
    const Array & lhs,
    const Array & rhs
)

Takes the elementwise maximum of two tensors.

Parameters:

• lhs Left operand.
• rhs Right operand.

Return: Elementwise maximum.

Strict elementwise maximum.

function Minimum¶

TENZR_APIArray Minimum(
    const Array & lhs,
    const Array & rhs
)

Takes the elementwise minimum of two tensors.

Parameters:

• lhs Left operand.
• rhs Right operand.

Return: Elementwise minimum.

Strict elementwise minimum.

function Remainder¶

TENZR_APIArray Remainder(
    const Array & lhs,
    const Array & rhs
)

Computes the elementwise remainder of lhs divided by rhs.

Parameters:

• lhs Dividend.
• rhs Divisor.

Return: lhs % rhs.

Strict elementwise remainder.

function Atan2¶

TENZR_APIArray Atan2(
    const Array & lhs,
    const Array & rhs
)

Computes the elementwise two-argument arctangent.

Parameters:

• lhs First operand.
• rhs Second operand.

Return: atan2(lhs, rhs).

Strict elementwise two-argument arctangent.

function Eq¶

TENZR_APIArray Eq(
    const Array & lhs,
    const Array & rhs
)

Compares two tensors for elementwise equality.

Parameters:

• lhs Left operand.
• rhs Right operand.

Return: Boolean array.

Strict elementwise equality.

function Ne¶

TENZR_APIArray Ne(
    const Array & lhs,
    const Array & rhs
)

Compares two tensors for elementwise inequality.

Parameters:

• lhs Left operand.
• rhs Right operand.

Return: Boolean array.

Strict elementwise inequality.

function Lt¶

TENZR_APIArray Lt(
    const Array & lhs,
    const Array & rhs
)

Tests whether lhs is elementwise less than rhs.

Parameters:

• lhs Left operand.
• rhs Right operand.

Return: Boolean array.

Strict elementwise less-than.

function Le¶

TENZR_APIArray Le(
    const Array & lhs,
    const Array & rhs
)

Tests whether lhs is elementwise less than or equal to rhs.

Parameters:

• lhs Left operand.
• rhs Right operand.

Return: Boolean array.

Strict elementwise less-than-or-equal.

function Gt¶

TENZR_APIArray Gt(
    const Array & lhs,
    const Array & rhs
)

Tests whether lhs is elementwise greater than rhs.

Parameters:

• lhs Left operand.
• rhs Right operand.

Return: Boolean array.

Strict elementwise greater-than.

function Ge¶

TENZR_APIArray Ge(
    const Array & lhs,
    const Array & rhs
)

Tests whether lhs is elementwise greater than or equal to rhs.

Parameters:

• lhs Left operand.
• rhs Right operand.

Return: Boolean array.

Strict elementwise greater-than-or-equal.

function Neg¶

TENZR_APIArray Neg(
    const Array & x
)

Negates each element.

Parameters:

• x Input array.

Return: -x.

Elementwise negation.

function Exp¶

TENZR_APIArray Exp(
    const Array & x
)

Computes the elementwise exponential.

Parameters:

• x Input array.

Return: exp(x).

Elementwise exponential.

function Log¶

TENZR_APIArray Log(
    const Array & x
)

Computes the elementwise natural logarithm.

Parameters:

• x Input array.

Return: log(x).

Elementwise natural logarithm.

function tanh¶

TENZR_APIArray tanh(
    const Array & x
)

Computes the elementwise hyperbolic tangent.

Parameters:

• x Input array.

Return: tanh(x).

Elementwise hyperbolic tangent.

function Abs¶

TENZR_APIArray Abs(
    const Array & x
)

Computes the elementwise absolute value.

Parameters:

• x Input array.

Return: |x|.

Elementwise absolute value.

function Sqrt¶

TENZR_APIArray Sqrt(
    const Array & x
)

Computes the elementwise square root.

Parameters:

• x Input array.

Return: sqrt(x).

Elementwise square root.

function Rsqrt¶

TENZR_APIArray Rsqrt(
    const Array & x
)

Computes the elementwise reciprocal square root.

Parameters:

• x Input array.

Return: 1 / sqrt(x).

Elementwise reciprocal square root.

function Sin¶

TENZR_APIArray Sin(
    const Array & x
)

Computes the elementwise sine.

Parameters:

• x Input array.

Return: sin(x).

Elementwise sine.

function Cos¶

TENZR_APIArray Cos(
    const Array & x
)

Computes the elementwise cosine.

Parameters:

• x Input array.

Return: cos(x).

Elementwise cosine.

function Tan¶

TENZR_APIArray Tan(
    const Array & x
)

Computes the elementwise tangent.

Parameters:

• x Input array.

Return: tan(x).

Elementwise tangent.

function Ceil¶

TENZR_APIArray Ceil(
    const Array & x
)

Rounds each element up to the nearest integer.

Parameters:

• x Input array.

Return: ceil(x).

Elementwise ceiling.

function Floor¶

TENZR_APIArray Floor(
    const Array & x
)

Rounds each element down to the nearest integer.

Parameters:

• x Input array.

Return: floor(x).

Elementwise floor.

function Round¶

TENZR_APIArray Round(
    const Array & x
)

Rounds each element to the nearest integer.

Parameters:

• x Input array.

Return: round(x).

Elementwise round-to-nearest.

function Sign¶

TENZR_APIArray Sign(
    const Array & x
)

Returns the elementwise sign (-1, 0, or 1).

Parameters:

• x Input array.

Return: sign(x).

Elementwise sign.

function ReduceSum¶

TENZR_APIArray ReduceSum(
    const Array & operand,
    const std::vector< int64_t > & axes ={},
    bool keepdims =false
)

Parameters:

• operand The tensor to reduce.
• axes Axes to reduce; reduces all axes when empty.
• keepdims Whether to retain reduced dimensions as size 1.
• operand Input array.
• axes Axes to reduce.
• keepdims Whether reduced axes are retained.

Return:

• The reduced tensor.
• The reduced array.

Sums operand over the given axes. Sum over axes; reduces all axes when axes is empty.

function ReduceMax¶

TENZR_APIArray ReduceMax(
    const Array & operand,
    const std::vector< int64_t > & axes ={},
    bool keepdims =false
)

Parameters:

• operand The tensor to reduce.
• axes Axes to reduce; reduces all axes when empty.
• keepdims Whether to retain reduced dimensions as size 1.
• operand Input array.
• axes Axes to reduce.
• keepdims Whether reduced axes are retained.

Return:

• The reduced tensor.
• The reduced array.

Computes the maximum of operand over the given axes. Maximum over axes; reduces all axes when axes is empty.

function ReduceMin¶

TENZR_APIArray ReduceMin(
    const Array & operand,
    const std::vector< int64_t > & axes ={},
    bool keepdims =false
)

Parameters:

• operand The tensor to reduce.
• axes Axes to reduce; reduces all axes when empty.
• keepdims Whether to retain reduced dimensions as size 1.
• operand Input array.
• axes Axes to reduce.
• keepdims Whether reduced axes are retained.

Return:

• The reduced tensor.
• The reduced array.

Computes the minimum of operand over the given axes. Minimum over axes; reduces all axes when axes is empty.

function ReduceProd¶

TENZR_APIArray ReduceProd(
    const Array & operand,
    const std::vector< int64_t > & axes ={},
    bool keepdims =false
)

Parameters:

• operand The tensor to reduce.
• axes Axes to reduce; reduces all axes when empty.
• keepdims Whether to retain reduced dimensions as size 1.
• operand Input array.
• axes Axes to reduce.
• keepdims Whether reduced axes are retained.

Return:

• The reduced tensor.
• The reduced array.

Multiplies operand over the given axes. Product over axes; reduces all axes when axes is empty.

function ReduceWindow¶

TENZR_APIArray ReduceWindow(
    const Array & operand,
    const Array & init_value,
    const std::vector< int64_t > & window_dimensions,
    const std::vector< int64_t > & window_strides ={},
    const std::vector< int64_t > & padding_low ={},
    const std::vector< int64_t > & padding_high ={},
    std::string_view reducer ="add"
)

Parameters:

• operand The input tensor.
• init_value The reduction identity value.
• window_dimensions The window size per dimension.
• window_strides The window stride per dimension.
• padding_low Padding added before each dimension.
• padding_high Padding added after each dimension.
• reducer The reduction name ("add", "max", "min", "mul").
• operand Input array.
• init_value Reduction identity.
• window_dimensions Window extent per dimension.
• window_strides Window stride per dimension; defaults to ones.
• padding_low Low padding per dimension; defaults to zeros.
• padding_high High padding per dimension; defaults to zeros.
• reducer One of "add", "max", "min", "multiply".

Return:

• The windowed reduction result.
• The windowed reduction.

Applies a windowed reduction with the named reducer over operand. Apply a windowed reduction with the given window geometry.

function slice¶

TENZR_APIArray slice(
    const Array & operand,
    const std::vector< int64_t > & start_indices,
    const std::vector< int64_t > & limit_indices,
    const std::vector< int64_t > & strides ={}
)

Parameters:

• operand The input tensor.
• start_indices Per-dimension start indices.
• limit_indices Per-dimension exclusive end indices.
• strides Per-dimension step sizes.
• operand Input array.
• start_indices Start index per dimension.
• limit_indices Exclusive end index per dimension.
• strides Step per dimension.

Return:

• The sliced tensor.
• The sliced array.

Extracts a strided slice of operand. Extract a strided region with explicit start/limit indices.

function Pad¶

TENZR_APIArray Pad(
    const Array & operand,
    const Array & padding_value,
    const std::vector< int64_t > & edge_padding_low,
    const std::vector< int64_t > & edge_padding_high,
    const std::vector< int64_t > & interior_padding ={}
)

Parameters:

• operand The input tensor.
• padding_value The value used for padding.
• edge_padding_low Padding added before each dimension.
• edge_padding_high Padding added after each dimension.
• interior_padding Padding inserted between elements per dimension.
• operand Input array.
• padding_value Scalar fill value; must match the operand dtype.
• edge_padding_low Low edge padding per dimension.
• edge_padding_high High edge padding per dimension.
• interior_padding Interior padding per dimension.

Return:

• The padded tensor.
• The padded array.

Pads operand with padding_value at the edges and between elements. Pad an array with edge and interior padding.

function Gather¶

TENZR_APIArray Gather(
    const Array & operand,
    const Array & start_indices,
    int64_t axis =0
)

Parameters:

• operand The tensor to gather from.
• start_indices The starting indices.
• axis The axis to gather along.
• operand Source array.
• start_indices 1-D index array.
• axis Axis to gather along.

Return:

• The gathered tensor.
• The gathered array with the index count prepended.

Gathers slices of operand at start_indices along axis. Gather 1-D slices of operand along axis at the given indices.

function DynamicSlice¶

TENZR_APIArray DynamicSlice(
    const Array & operand,
    const Array & start_indices,
    const std::vector< int64_t > & slice_sizes
)

Parameters:

• operand The tensor to slice.
• start_indices The starting offset.
• slice_sizes The size of each sliced dimension.
• operand Source array.
• start_indices Vector holding one start index per dimension.
• slice_sizes Static slice sizes.

Return:

• The sliced tensor.
• The sliced array.

Extracts a dynamic slice of fixed slice_sizes starting at start_indices. Dynamically slice, unpacking a single vector of start indices.

function DynamicSlice¶

TENZR_APIArray DynamicSlice(
    const Array & operand,
    const std::vector< Array > & start_indices,
    const Shape & slice_sizes
)

Parameters:

• operand The tensor to slice.
• start_indices Per-dimension starting offsets.
• slice_sizes The size of each sliced dimension.
• operand Source array.
• start_indices One start index per dimension.
• slice_sizes Static slice shape.

Return:

• The sliced tensor.
• The sliced array.

Extracts a dynamic slice with one start index array per dimension. Dynamically slice with a per-dimension list of start indices.

function DynamicUpdateSlice¶

TENZR_APIArray DynamicUpdateSlice(
    const Array & operand,
    const Array & update,
    const Array & start_indices
)

Parameters:

• operand The tensor to update.
• update The replacement values.
• start_indices The starting offset.
• operand Target array.
• update Values to write.
• start_indices Vector holding one start index per dimension.

Return:

• The updated tensor.
• The updated array.

Replaces a slice of operand at start_indices with update. Dynamically update a slice, unpacking a single vector of start indices.

function DynamicUpdateSlice¶

TENZR_APIArray DynamicUpdateSlice(
    const Array & operand,
    const Array & update,
    const std::vector< Array > & start_indices
)

Parameters:

• operand The tensor to update.
• update The replacement values.
• start_indices Per-dimension starting offsets.
• operand Target array.
• update Values to write.
• start_indices One start index per dimension.

Return:

• The updated tensor.
• The updated array.

Replaces a slice of operand starting at per-dimension indices with update. Dynamically write update into operand at per-dimension start indices.

function ScatterAdd¶

TENZR_APIArray ScatterAdd(
    const Array & operand,
    const Array & indices,
    const Array & updates,
    int64_t axis =0
)

Parameters:

• operand The tensor to update.
• indices The scatter indices.
• updates The values to add.
• axis The axis to scatter along.
• operand Target array.
• indices Index array.
• updates Values to add.
• axis Axis to scatter along.

Return:

• The updated tensor.
• The scattered array.

Scatters and adds updates into operand at indices. Scatter updates with addition.

function Scatter¶

TENZR_APIArray Scatter(
    const Array & operand,
    const Array & indices,
    const Array & updates,
    int64_t axis =0,
    std::string_view update ="add"
)

Parameters:

• operand The tensor to update.
• indices The scatter indices.
• updates The values to scatter.
• axis The axis to scatter along.
• update The update kind ("add", "set", "mul", ...).
• operand Target array.
• indices 1-D index array.
• updates Values to scatter.
• axis Axis to scatter along.
• update Combination mode: "add", "set", "multiply", "min", or "max".

Return:

• The updated tensor.
• The scattered array.

Scatters updates into operand at indices using the named update operation. Scatter updates into operand along axis with a combining operator.

function Argsort¶

TENZR_APIArray Argsort(
    const Array & operand,
    int64_t axis =-1,
    bool descending =false
)

Parameters:

• operand The tensor to inspect.
• axis The axis to sort along.
• descending Whether to sort in descending order.
• operand Input array.
• axis Sort axis.
• descending Sort order.

Return:

• The sort permutation.
• The sort permutation.

Returns the indices that would sort operand along axis. Indices that sort operand along axis.

function Sort¶

TENZR_APIArray Sort(
    const Array & operand,
    int64_t axis =-1,
    bool descending =false
)

Parameters:

• operand The tensor to sort.
• axis The axis to sort along.
• descending Whether to sort in descending order.
• operand Input array.
• axis Sort axis.
• descending Sort order.

Return:

• The sorted tensor.
• The sorted values.

Sorts operand along axis. Sort operand along axis.

function Argmax¶

TENZR_APIArray Argmax(
    const Array & operand,
    int64_t axis =-1
)

Parameters:

• operand The tensor to inspect.
• axis The axis to reduce.
• operand Input array.
• axis Reduction axis.

Return:

• The argmax indices.
• Int32 index array.

Returns the indices of the maximum values along axis. Index of the maximum along axis, choosing the lowest index on ties.

function Argmin¶

TENZR_APIArray Argmin(
    const Array & operand,
    int64_t axis =-1
)

Parameters:

• operand The tensor to inspect.
• axis The axis to reduce.
• operand Input array.
• axis Reduction axis.

Return:

• The argmin indices.
• Int32 index array.

Returns the indices of the minimum values along axis. Index of the minimum along axis, choosing the lowest index on ties.

function Concatenate¶

TENZR_APIArray Concatenate(
    const std::vector< Array > & arrays,
    int64_t axis =0
)

Parameters:

• arrays The tensors to concatenate.
• axis The axis along which to concatenate.
• arrays Arrays to join.
• axis Concatenation axis.

Return:

• The concatenated tensor.
• The concatenated array.

Concatenates arrays along an existing axis. Join arrays along an existing axis; dtypes and ranks must match.

function Stack¶

TENZR_APIArray Stack(
    const std::vector< Array > & arrays,
    int64_t axis =0
)

Parameters:

• arrays The tensors to stack.
• axis Where to insert the new axis.
• arrays Arrays to stack; all must share a shape.
• axis Position of the new axis.

Return:

• The stacked tensor.
• The stacked array.

Stacks arrays along a new axis. Join arrays along a new axis.

function Conv2d¶

TENZR_APIArray Conv2d(
    const Array & input,
    const Array & weight,
    std::pair< int64_t, int64_t > stride ={1, 1},
    std::pair< int64_t, int64_t > padding ={0, 0}
)

Parameters:

• input The input tensor.
• weight The convolution kernel.
• stride The vertical and horizontal stride.
• padding The vertical and horizontal padding.
• input Input array of shape [B, C, H, W].
• weight Kernel of shape [OC, C, KH, KW].
• stride Stride height and width.
• padding Padding height and width.

Return:

• The convolution result.
• The convolved array.

Computes a 2D cross-correlation convolution. NCHW/OIHW 2-D convolution with symmetric padding.

function ConvGeneral¶

TENZR_APIArray ConvGeneral(
    const Array & lhs,
    const Array & rhs,
    const std::vector< int64_t > & window_strides,
    const std::vector< int64_t > & lhs_dilation,
    const std::vector< int64_t > & rhs_dilation,
    const std::vector< int64_t > & padding_low,
    const std::vector< int64_t > & padding_high,
    const std::vector< int64_t > & lhs_spec ={0, 1, 2, 3},
    const std::vector< int64_t > & rhs_spec ={0, 1, 2, 3},
    const std::vector< int64_t > & out_spec ={0, 1, 2, 3}
)

Parameters:

• lhs The left-hand (input) tensor.
• rhs The right-hand (kernel) tensor.
• window_strides The stride per spatial dimension.
• lhs_dilation Dilation of the input per dimension.
• rhs_dilation Dilation of the kernel per dimension.
• padding_low Low padding per dimension.
• padding_high High padding per dimension.
• lhs_spec Dimension numbers of the input.
• rhs_spec Dimension numbers of the kernel.
• out_spec Dimension numbers of the output.
• lhs Input operand.
• rhs Kernel operand.
• window_strides Window strides.
• lhs_dilation Input dilation.
• rhs_dilation Kernel dilation.
• padding_low Low edge padding.
• padding_high High edge padding.
• lhs_spec Dimension layout of lhs.
• rhs_spec Dimension layout of rhs.
• out_spec Dimension layout of the output.

Return:

• The convolution result.
• The convolved array.

Computes a general N-D convolution with dilations, padding, and dimension numbers. General N-D convolution with explicit dimension specs and padding.

function TakeAlongAxis¶

TENZR_APIArray TakeAlongAxis(
    const Array & operand,
    const Array & indices,
    int64_t axis =-1
)

Parameters:

• operand The tensor to gather from.
• indices The indices with the same rank as operand.
• axis The axis to gather along.
• operand Input array.
• indices Index array, broadcast with the operand.
• axis Axis along which indices address the operand.

Return:

• The gathered tensor.
• The gathered array.

Gathers values along axis using matching indices. Gather operand entries selected by a same-rank index tensor.

function PutAlongAxis¶

TENZR_APIArray PutAlongAxis(
    const Array & operand,
    const Array & indices,
    const Array & values,
    int64_t axis =-1,
    std::string_view update ="set"
)

Parameters:

• operand The tensor to update.
• indices The indices with the same rank as operand.
• values The values to place.
• axis The axis to update along.
• update The update kind ("set", "add", "mul", ...).
• operand Target array.
• indices Index array.
• values Values to write; must match the operand dtype.
• axis Axis along which indices address the operand.
• update Combination mode.

Return:

• The updated tensor.
• The updated array.

Updates operand along axis at indices using values and the named update. Scatter values into operand at positions selected by indices.

function TopK¶

TENZR_APITopKResult TopK(
    const Array & operand,
    int64_t k,
    int64_t axis =-1,
    bool largest =true
)

Parameters:

• operand The tensor to inspect.
• k The number of elements to return.
• axis The axis to search along.
• largest Whether to select the largest (or smallest) values.
• operand Input array.
• k Number of elements to select.
• axis Selection axis.
• largest Select the largest when true, the smallest otherwise.

Return:

• A pair of values and indices.
• Pair of {values, indices}.

Returns the top k values and their indices along axis. Return the k largest or smallest values and their indices along axis.

function Reverse¶

TENZR_APIArray Reverse(
    const Array & operand,
    const std::vector< int64_t > & dimensions
)

Parameters:

• operand The tensor to reverse.
• dimensions The dimensions to reverse.
• operand Input array.
• dimensions Axes to reverse.

Return:

• The reversed tensor.
• The reversed array.

Reverses operand along the given dimensions. Reverse the given dimensions.

function DotGeneral¶

TENZR_APIArray DotGeneral(
    const Array & lhs,
    const Array & rhs,
    const std::vector< int64_t > & lhs_contracting_dims,
    const std::vector< int64_t > & rhs_contracting_dims,
    const std::vector< int64_t > & lhs_batch_dims ={},
    const std::vector< int64_t > & rhs_batch_dims ={}
)

Parameters:

• lhs The left-hand tensor.
• rhs The right-hand tensor.
• lhs_contracting_dims Contracted dimensions of lhs.
• rhs_contracting_dims Contracted dimensions of rhs.
• lhs_batch_dims Batch dimensions of lhs.
• rhs_batch_dims Batch dimensions of rhs.
• lhs Left operand.
• rhs Right operand.
• lhs_contracting_dims Contracted dimensions of lhs.
• rhs_contracting_dims Contracted dimensions of rhs.
• lhs_batch_dims Batch dimensions of lhs.
• rhs_batch_dims Batch dimensions of rhs.

Return:

• The contraction result.
• The contraction result.

Computes a general contracted tensor product with optional batch dimensions. General tensor contraction with explicit contracting and batch dimensions.

function BroadcastInDim¶

TENZR_APIArray BroadcastInDim(
    const Array & operand,
    const Shape & target_shape,
    const std::vector< int64_t > & broadcast_dimensions
)

Parameters:

• operand The tensor to broadcast.
• target_shape The output shape.
• broadcast_dimensions Target dimension for each source dimension.
• operand Input array.
• target_shape Shape to broadcast to.
• broadcast_dimensions Source axis alignment; currently unused.

Return:

• The broadcasted tensor.
• The broadcast array.

Broadcasts operand to target_shape by mapping source dims through broadcast_dimensions. Broadcast operand to target_shape.

function Reshape¶

TENZR_APIArray Reshape(
    const Array & operand,
    const Shape & new_shape
)

Parameters:

• operand The tensor to reshape.
• new_shape The output shape.
• operand Input array.
• new_shape Target shape.

Return:

• The reshaped tensor.
• The reshaped array.

Returns operand reshaped to new_shape. Reshape without changing the element count.

function Transpose¶

TENZR_APIArray Transpose(
    const Array & operand,
    const std::vector< int64_t > & permutation
)

Parameters:

• operand The tensor to transpose.
• permutation The permutation of dimensions.
• operand Input array.
• permutation Target position of each source dimension.

Return:

• The transposed tensor.
• The transposed array.

Permutes the dimensions of operand according to permutation. Permute dimensions.

function ConvertElementType¶

TENZR_APIArray ConvertElementType(
    const Array & operand,
    DType new_dtype
)

Parameters:

• operand The tensor to convert.
• new_dtype The target dtype.
• operand Input array.
• new_dtype Destination dtype.

Return:

• The converted tensor.
• The converted array.

Casts operand to new_dtype. Convert elements to a new dtype.

function Select¶

TENZR_APIArray Select(
    const Array & pred,
    const Array & on_true,
    const Array & on_false
)

Parameters:

• pred The boolean predicate.
• on_true Values selected where pred is true.
• on_false Values selected where pred is false.
• pred Boolean predicate.
• on_true Value where pred is true.
• on_false Value where pred is false.

Return:

• The selected tensor.
• The selected array.

Selects on_true or on_false per element according to pred. Choose between two branches elementwise.

function Clamp¶

TENZR_APIArray Clamp(
    const Array & min_val,
    const Array & operand,
    const Array & max_val
)

Parameters:

• min_val The lower bound.
• operand The tensor to clamp.
• max_val The upper bound.
• min_val Lower bound.
• operand Input array.
• max_val Upper bound.

Return:

• The clamped tensor.
• The clamped array.

Clamps operand to the closed interval [min_val, max_val]. Clamp operand into [min_val, max_val].

function AllReduce¶

TENZR_APIArray AllReduce(
    const Array & x,
    std::string_view op ="sum",
    const std::vector< std::vector< int64_t > > & replica_groups ={}
)

Parameters:

• x The tensor to reduce.
• op The reduction name ("sum", "max", "min", ...).
• replica_groups Groups of replica ids participating.
• x Input array.
• op Reduction op, e.g. "sum".
• replica_groups Device groups to reduce over.

Return:

• The reduced tensor.
• The reduced array.

Reduces x across replicas with the named op. All-reduce across replica groups; a local identity copy on a single replica.

function AllGather¶

TENZR_APIArray AllGather(
    const Array & x,
    int64_t all_gather_dim =0,
    const std::vector< std::vector< int64_t > > & replica_groups ={}
)

Parameters:

• x The tensor to gather.
• all_gather_dim The dimension along which to concatenate.
• replica_groups Groups of replica ids participating.
• x Input array.
• all_gather_dim Dimension to concatenate gathered shards along.
• replica_groups Device groups involved.

Return:

• The gathered tensor.
• The gathered array.

Gathers x across replicas along all_gather_dim. All-gather along all_gather_dim; a local copy on a single replica.

function AllToAll¶

TENZR_APIArray AllToAll(
    const Array & x,
    int64_t split_dim =0,
    int64_t concat_dim =0,
    int64_t split_count =1,
    const std::vector< std::vector< int64_t > > & replica_groups ={}
)

Parameters:

• x The tensor to exchange.
• split_dim The dimension to split.
• concat_dim The dimension to concatenate.
• split_count The number of splits per replica.
• replica_groups Groups of replica ids participating.
• x Input array.
• split_dim Dimension divided among participants.
• concat_dim Dimension where exchanged chunks are concatenated.
• split_count Number of participants.
• replica_groups Device groups involved.

Return:

• The exchanged tensor.
• The exchanged array.

Performs an all-to-all split and concatenation across replicas. All-to-all exchange splitting split_dim and concatenating along concat_dim.

function ReduceScatter¶

TENZR_APIArray ReduceScatter(
    const Array & x,
    std::string_view op ="sum",
    int64_t scatter_dim =0,
    const std::vector< std::vector< int64_t > > & replica_groups ={}
)

Parameters:

• x The tensor to reduce and scatter.
• op The reduction name ("sum", "max", "min", ...).
• scatter_dim The dimension along which to scatter.
• replica_groups Groups of replica ids participating.
• x Input array.
• op Reduction op, e.g. "sum".
• scatter_dim Dimension divided after reduction.
• replica_groups Device groups involved.

Return:

• The reduced and scattered tensor.
• The reduced-and-scattered array.

Reduces x across replicas and scatters the result along scatter_dim. Reduce then scatter a dimension across replicas.

function Psum¶

TENZR_APIArray Psum(
    const Array & x,
    std::string_view axis_name =""
)

Parameters:

• x The tensor to reduce.
• axis_name The parallel axis name.
• x Input array.

Return:

• The summed tensor.
• The summed array.

Sums x across the named parallel axis. Sum along a named mesh axis.

function Pmean¶

TENZR_APIArray Pmean(
    const Array & x,
    std::string_view axis_name =""
)

Parameters:

• x The tensor to reduce.
• axis_name The parallel axis name.
• x Input array.

Return:

• The averaged tensor.
• The averaged array.

Averages x across the named parallel axis. Mean along a named mesh axis.

function WithShardingConstraint¶

TENZR_APIArray WithShardingConstraint(
    const Array & x,
    const NamedSharding & sharding
)

Parameters:

• x The tensor to constrain.
• sharding The desired named sharding.
• x Input array.
• sharding Named sharding to constrain to.

Return:

• The constrained tensor.
• The constrained array.

Constrains x to sharding and returns it. Attach a sharding constraint to x.