apn_mojo.batch¶
This package provides multidimensional batches, boolean masks, vmap, and
NumPy-style array functions. Copies and views keep their values across
updates. Start with the
batch tutorial and the
vectorization tutorial; the full lift
contract is described in apn_mojo.
Batches of any rank, Boolean masks, vmap, and NumPy-style functions of batches.
A Batch[T] holds Integers, Rationals, Floats, Complex numbers or balls with
value semantics: selections share values, and updates publish new storage, so
no batch sees another's later update. Long operations and mappings run on a
worker pool with the same results and errors as a sequential loop.
from apn_mojo import batch gives the NumPy-style namespace: batch.exp(xs),
batch.atan2(ys, xs), batch.where(xs > 0, xs, zero), batch.zeros[Float](n),
batch.cumsum(xs). Each elementwise function maps the scalar function of the
same name over every element, with broadcasting; the scalar functions stay at
the package root.
Summary¶
| Name | Kind | Summary |
|---|---|---|
abs |
function | The absolute value of every element, like numpy.abs. |
acos |
function | The arccosine of every element, like numpy.arccos. |
acosh |
function | The inverse hyperbolic cosine of every element, like numpy.arccosh. |
add |
function | a + b elementwise, like numpy.add. |
angle |
function | The argument of every element in (-pi, pi], like numpy.angle: Floats rounded once, or balls. |
arange |
function | start, start + step, ... up to but excluding stop, like numpy.arange. |
argmax |
function | The row-major index of the first greatest element (numpy's argmax); of the first NaN in a Float batch. |
argmin |
function | The row-major index of the first least element (numpy's argmin); of the first NaN in a Float batch. |
asin |
function | The arcsine of every element, like numpy.arcsin. |
asinh |
function | The inverse hyperbolic sine of every element, like numpy.arcsinh. |
atan |
function | The arctangent of every element, like numpy.arctan. |
atan2 |
function | The angle of each point (x, y), like numpy.arctan2. |
atanh |
function | The inverse hyperbolic tangent of every element, like numpy.arctanh. |
beta |
function | The beta function of each pair, like scipy.special.beta. |
betainc |
function | The regularized incomplete beta function, like scipy.special.betainc. |
betaln |
function | log|beta(a, b)| of each pair, like scipy.special.betaln. |
ceil |
function | The smallest Integer at least each element, exactly, like numpy.ceil. |
clip |
function | Each value limited to [a_min, a_max], like numpy.clip. |
comb |
function | The number of ways to choose k of N things, exactly, like scipy.special.comb with exact=True; 0 when k > N, N < 0 or k < 0. |
concatenate |
function | The batches joined along an existing axis, like numpy.concatenate. |
conjugate |
function | The complex conjugate of every element, like numpy.conjugate. |
cos |
function | The cosine of every element, like numpy.cos. |
cosh |
function | The hyperbolic cosine of every element, like numpy.cosh. |
cumprod |
function | Running products along an axis, like numpy.cumprod; without an axis, over the batch flattened in row-major order. |
cumsum |
function | Running sums along an axis, like numpy.cumsum; without an axis, over the batch flattened in row-major order. |
digamma |
function | The digamma function of every real element or ball, like scipy.special.digamma. |
divide |
function | a / b elementwise, like numpy.divide: Integers give exact Rationals; other families as for add. |
dot |
function | The sum of pairwise products, exactly. |
erf |
function | The error function of every real element or ball, like scipy.special.erf. |
erfc |
function | The complementary error function of every real element or ball, like scipy.special.erfc. |
erfi |
function | The imaginary error function of every real element or ball, like scipy.special.erfi. |
erfinv |
function | The inverse error function of every real element or ball, like scipy.special.erfinv. |
exp |
function | apn_mojo.exp of every element, like numpy.exp. |
exp2 |
function | 2**x of every real element or ball, like numpy.exp2. |
expi |
function | The exponential integral Ei of every real element or ball, like scipy.special.expi. |
expm1 |
function | exp(x) - 1 of every real element or ball, like numpy.expm1. |
floor |
function | The largest Integer at most each element, exactly, like numpy.floor but with Integer results. |
fresnel |
function | The Fresnel integrals (S, C) of every real element or ball, like scipy.special.fresnel. |
full |
function | A batch holding value everywhere, like numpy.full; the family and format are the value's. |
gamma |
function | The gamma function of every real element or ball, like scipy.special.gamma. |
gammainc |
function | The regularized lower incomplete gamma function, like scipy.special.gammainc. |
gammaincc |
function | The regularized upper incomplete gamma function, like scipy.special.gammaincc. |
gammaln |
function | log|gamma(x)| of every real element or ball, like scipy.special.gammaln. |
hyp1f1 |
function | Kummer's confluent hypergeometric function, like scipy.special.hyp1f1. |
hyp2f1 |
function | The Gauss hypergeometric function, like scipy.special.hyp2f1. |
imag |
function | The imaginary part of every element, like numpy.imag: Floats or balls. |
lambertw |
function | Branch k of the Lambert W function of every real element or ball, like scipy.special.lambertw. |
linspace |
function | num evenly spaced values from start to stop, like numpy.linspace. |
log |
function | The natural logarithm of every element, like numpy.log; Complex elements take the principal branch. |
log10 |
function | The base-10 logarithm of every real element or ball, like numpy.log10. |
log1p |
function | log(1 + x) of every real element or ball, like numpy.log1p. |
log2 |
function | The base-2 logarithm of every real element or ball, like numpy.log2. |
log_ndtr |
function | The logarithm of ndtr of every real element or ball, like scipy.special.log_ndtr. |
max |
function | The greatest element (numpy's max). |
maximum |
function | The larger of each pair, like numpy.maximum; a NaN wins, as in numpy. |
min |
function | The least element (numpy's min). |
minimum |
function | The smaller of each pair, like numpy.minimum; as for maximum. |
multiply |
function | a * b elementwise, like numpy.multiply; families as for add. |
ndtr |
function | The standard normal distribution function of every real element or ball, like scipy.special.ndtr. |
ndtri |
function | The inverse of ndtr of every real element or ball, like scipy.special.ndtri. |
ones |
function | A batch of ones, like numpy.ones. |
poch |
function | The Pochhammer symbol (z)_m of each pair, like scipy.special.poch. |
polygamma |
function | The nth derivative of digamma at each x, like scipy.special.polygamma; n is Integers. |
pow |
function | base ** exponent elementwise, like numpy.power, for real, Complex and ball values. |
pow_int |
function | Each value to an Integer power, rounded once. |
prod |
function | The exact product of every element. |
real |
function | The real part of every element, like numpy.real: Floats or balls. |
reciprocal |
function | 1 / x elementwise, like numpy.reciprocal: Integers give exact Rationals, Rationals stay exact, the other families keep theirs. |
rootn |
function | The real nth root of every real element or ball. |
round |
function | Each element rounded to the nearest Integer, ties to even, like numpy.round. |
shichi |
function | The hyperbolic sine and cosine integrals of every real element or ball, like scipy.special.shichi. |
sici |
function | The sine and cosine integrals of every real element or ball, like scipy.special.sici. |
sin |
function | The sine of every element, like numpy.sin. |
sin_cos |
function | The sines and the cosines of every real element or ball, computed together. |
sinh |
function | The hyperbolic sine of every element, like numpy.sinh. |
sqrt |
function | The square root of every element, like numpy.sqrt. |
stack |
function | The batches joined along a new axis, like numpy.stack. |
subtract |
function | a - b elementwise, like numpy.subtract; families as for add. |
sum |
function | Sum every element exactly. |
tan |
function | The tangent of every element, like numpy.tan. |
tanh |
function | The hyperbolic tangent of every element, like numpy.tanh. |
trunc |
function | Each element rounded toward zero to an Integer, like numpy.trunc. |
vdot |
function | The dot product with the first operand conjugated. |
vmap |
function | Map a scalar function over batches. |
where |
function | Elements of a where condition holds and of b elsewhere, like numpy.where. |
zeros |
function | A batch of zeros, like numpy.zeros. |
zeta |
function | The Hurwitz zeta function zeta(x, q) of each pair, like scipy.special.zeta. |
Batch |
struct | A batch of numbers of any rank, with value semantics. |
Mask |
struct | An immutable batch of Booleans, returned by batch comparisons. |
Functions¶
abs¶
Source: apn_mojo/batch/functions.mojo
def abs[T: ImplicitlyCopyable & Deinitable, //](
values: Batch[T],
*,
context: _MagnitudeContext[T] = _MagnitudeContext[T](),
) raises -> Batch[_Magnitude[T]]
The absolute value of every element, like numpy.abs.
Integers, Rationals and Floats keep their family exactly; Complex numbers give their magnitudes as Floats rounded once, and balls give balls.
Parameters
T(ImplicitlyCopyable & Deinitable, inferred): The element type, inferred.
Arguments
values(Batch[T]): The batch.context(_MagnitudeContext[T]): The rounding of Complex magnitudes, or the ball context; none for the other families.
Returns
Batch[...]: A batch of the same shape.
Raises
Error: As the scalar function does, with the failing element's position.
acos¶
Source: apn_mojo/batch/functions.mojo
def acos[T: ImplicitlyCopyable & Deinitable, //](
values: Batch[T],
*,
context: _ContextOf[T] = _ContextOf[T](),
) raises -> Batch[_Mapped[T]]
The arccosine of every element, like numpy.arccos.
Parameters
T(ImplicitlyCopyable & Deinitable, inferred): The element type, inferred.
Arguments
values(Batch[T]): The batch.context(_ContextOf[T]): The family's context, as for the scalar function; exact elements need one.
Returns
Batch[...]: A batch of the same shape.
Raises
Error: As the scalar function does, with the failing element's position.
acosh¶
Source: apn_mojo/batch/functions.mojo
def acosh[T: ImplicitlyCopyable & Deinitable, //](
values: Batch[T],
*,
context: _ContextOf[T] = _ContextOf[T](),
) raises -> Batch[_Mapped[T]]
The inverse hyperbolic cosine of every element, like numpy.arccosh.
Parameters
T(ImplicitlyCopyable & Deinitable, inferred): The element type, inferred.
Arguments
values(Batch[T]): The batch.context(_ContextOf[T]): The family's context, as for the scalar function; exact elements need one.
Returns
Batch[...]: A batch of the same shape.
Raises
Error: As the scalar function does, with the failing element's position.
add¶
Source: apn_mojo/batch/functions.mojo
def add[
A: _MapArgument & ImplicitlyCopyable & Deinitable, B: _MapArgument &
ImplicitlyCopyable & Deinitable, //,
](
a: A,
b: B,
) raises -> Batch[_Arithmetic[A, B, Integer]]
a + b elementwise, like numpy.add.
Integers and Rationals stay exact; a Float value gives Floats rounded once to the operands' format, a Complex value Complex numbers, a ball balls.
Parameters
A(inferred): The type ofa, inferred.B(inferred): The type ofb, inferred.
Arguments
a(A): A batch or a scalar.b(B): A batch or a scalar; at least one argument is a batch.
Returns
Batch[...]: A batch of the broadcast shape.
Raises
Error: When the shapes do not broadcast, or as the scalar function does.
1 more overload
`a + b` elementwise, each sum rounded once with `context`: an ArithmeticContext gives Floats (Complex for complex values), a ComplexContext Complex numbers, a BallContext balls.
def add[
A: _MapArgument & ImplicitlyCopyable & Deinitable, B: _MapArgument &
ImplicitlyCopyable & Deinitable, C: ImplicitlyCopyable, //,
](
a: A,
b: B,
*,
context: C,
) raises -> Batch[_WithContext[_Arithmetic[A, B, Integer], C]]angle¶
Source: apn_mojo/batch/functions.mojo
def angle[T: ImplicitlyCopyable & Deinitable, //](
values: Batch[T],
*,
context: _MagnitudeContext[T] = _MagnitudeContext[T](),
) raises -> Batch[_Magnitude[T]]
The argument of every element in (-pi, pi], like numpy.angle: Floats rounded once, or balls.
Parameters
T(ImplicitlyCopyable & Deinitable, inferred): The element type, inferred.
Arguments
values(Batch[T]): The batch.context(_MagnitudeContext[T]): The family's context, as for the scalar function; exact elements need one.
Returns
Batch[...]: A batch of the same shape.
Raises
Error: As the scalar function does, with the failing element's position.
arange¶
Source: apn_mojo/batch/creation.mojo
def arange[T: ImplicitlyCopyable & Deinitable](
start: Rational,
stop: Rational,
step: Rational = Rational(1),
*,
context: _FormatContext[T] = _FormatContext[T](),
) raises -> Batch[T]
start, start + step, ... up to but excluding stop, like numpy.arange.
Each element is computed exactly: Integer elements must be whole, and a
Float element is start + k * step rounded once, so no rounding error
accumulates.
Parameters
T(ImplicitlyCopyable & Deinitable): The element type: Integer, Rational or Float.
Arguments
start(Rational): The first value.stop(Rational): The bound, excluded.step(Rational): The step, positive or negative.context(_FormatContext[T]): The Float format; 128 bits by default.
Returns
Batch[T]: A vector of ceil((stop - start) / step) elements, or none.
Raises
Error: For a zero step, or a value that is not whole for Integer elements.
1 more overload
`0, 1, ...` below `stop`, like `numpy.arange(stop)`.
def arange[T: ImplicitlyCopyable & Deinitable](
stop: Rational,
*,
context: _FormatContext[T] = _FormatContext[T](),
) raises -> Batch[T]argmax¶
Source: apn_mojo/batch/reductions.mojo
def argmax[T: ImplicitlyCopyable](values: T) raises -> Integer
The row-major index of the first greatest element (numpy's argmax); of the first NaN in a Float batch.
Parameters
T(ImplicitlyCopyable): The batch type.
Arguments
values(T): A batch, or any selection of one.
Returns
Integer: The flat index.
Raises
Error: When the batch is empty.
1 more overload
The index along `axis` of the first greatest element of each lane.
def argmax[T: ImplicitlyCopyable](
values: T,
*,
axis: Int,
keepdims: Bool = False,
) raises -> Batch[Integer] where conforms_to(T, _BatchShape)argmin¶
Source: apn_mojo/batch/reductions.mojo
def argmin[T: ImplicitlyCopyable](values: T) raises -> Integer
The row-major index of the first least element (numpy's argmin); of the first NaN in a Float batch.
Parameters
T(ImplicitlyCopyable): The batch type.
Arguments
values(T): A batch, or any selection of one.
Returns
Integer: The flat index.
Raises
Error: When the batch is empty.
1 more overload
The index along `axis` of the first least element of each lane.
def argmin[T: ImplicitlyCopyable](
values: T,
*,
axis: Int,
keepdims: Bool = False,
) raises -> Batch[Integer] where conforms_to(T, _BatchShape)asin¶
Source: apn_mojo/batch/functions.mojo
def asin[T: ImplicitlyCopyable & Deinitable, //](
values: Batch[T],
*,
context: _ContextOf[T] = _ContextOf[T](),
) raises -> Batch[_Mapped[T]]
The arcsine of every element, like numpy.arcsin.
Parameters
T(ImplicitlyCopyable & Deinitable, inferred): The element type, inferred.
Arguments
values(Batch[T]): The batch.context(_ContextOf[T]): The family's context, as for the scalar function; exact elements need one.
Returns
Batch[...]: A batch of the same shape.
Raises
Error: As the scalar function does, with the failing element's position.
asinh¶
Source: apn_mojo/batch/functions.mojo
def asinh[T: ImplicitlyCopyable & Deinitable, //](
values: Batch[T],
*,
context: _ContextOf[T] = _ContextOf[T](),
) raises -> Batch[_Mapped[T]]
The inverse hyperbolic sine of every element, like numpy.arcsinh.
Parameters
T(ImplicitlyCopyable & Deinitable, inferred): The element type, inferred.
Arguments
values(Batch[T]): The batch.context(_ContextOf[T]): The family's context, as for the scalar function; exact elements need one.
Returns
Batch[...]: A batch of the same shape.
Raises
Error: As the scalar function does, with the failing element's position.
atan¶
Source: apn_mojo/batch/functions.mojo
def atan[T: ImplicitlyCopyable & Deinitable, //](
values: Batch[T],
*,
context: _ContextOf[T] = _ContextOf[T](),
) raises -> Batch[_Mapped[T]]
The arctangent of every element, like numpy.arctan.
Parameters
T(ImplicitlyCopyable & Deinitable, inferred): The element type, inferred.
Arguments
values(Batch[T]): The batch.context(_ContextOf[T]): The family's context, as for the scalar function; exact elements need one.
Returns
Batch[...]: A batch of the same shape.
Raises
Error: As the scalar function does, with the failing element's position.
atan2¶
Source: apn_mojo/batch/functions.mojo
def atan2[
A: _MapArgument & ImplicitlyCopyable & Deinitable, B: _MapArgument &
ImplicitlyCopyable & Deinitable, //,
](
y: A,
x: B,
*,
context: _ContextOf[_Pair[A, B]] = _ContextOf[_Pair[A, B]](),
) raises -> Batch[_Pair[A, B]]
The angle of each point (x, y), like numpy.arctan2.
Two-argument functions broadcast their arguments; either may be a scalar.
Parameters
A(inferred): The type ofy, inferred.B(inferred): The type ofx, inferred.
Arguments
y(A): A batch or a scalar.x(B): A batch or a scalar; at least one argument is a batch.context(_ContextOf[_Pair[A, B]]): The family's context; exact values need one.
Returns
Batch[...]: A batch of the broadcast shape: Floats, or balls if either argument holds balls.
Raises
Error: When the shapes do not broadcast, or as the scalar function does.
atanh¶
Source: apn_mojo/batch/functions.mojo
def atanh[T: ImplicitlyCopyable & Deinitable, //](
values: Batch[T],
*,
context: _ContextOf[T] = _ContextOf[T](),
) raises -> Batch[_Mapped[T]]
The inverse hyperbolic tangent of every element, like numpy.arctanh.
Parameters
T(ImplicitlyCopyable & Deinitable, inferred): The element type, inferred.
Arguments
values(Batch[T]): The batch.context(_ContextOf[T]): The family's context, as for the scalar function; exact elements need one.
Returns
Batch[...]: A batch of the same shape.
Raises
Error: As the scalar function does, with the failing element's position.
beta¶
Source: apn_mojo/batch/functions.mojo
def beta[
A: _MapArgument & ImplicitlyCopyable & Deinitable, B: _MapArgument &
ImplicitlyCopyable & Deinitable, //,
](
a: A,
b: B,
*,
context: _ContextOf[_Pair[A, B]] = _ContextOf[_Pair[A, B]](),
) raises -> Batch[_Pair[A, B]]
The beta function of each pair, like scipy.special.beta.
Parameters
A(inferred): The type ofa, inferred.B(inferred): The type ofb, inferred.
Arguments
a(A): A batch or a scalar.b(B): A batch or a scalar; at least one argument is a batch.context(_ContextOf[_Pair[A, B]]): The family's context; exact values need one.
Returns
Batch[...]: A batch of the broadcast shape.
Raises
Error: When the shapes do not broadcast, or as the scalar function does.
betainc¶
Source: apn_mojo/batch/functions.mojo
def betainc[
A: _MapArgument & ImplicitlyCopyable & Deinitable, B: _MapArgument &
ImplicitlyCopyable & Deinitable, D: _MapArgument & ImplicitlyCopyable &
Deinitable, //,
](
a: A,
b: B,
x: D,
*,
context: _ContextOf[_Triple[A, B, D]] = _ContextOf[_Triple[A, B, D]](),
) raises -> Batch[_Triple[A, B, D]]
The regularized incomplete beta function, like scipy.special.betainc.
Parameters
A(inferred): The type ofa, inferred.B(inferred): The type ofb, inferred.D(inferred): The type ofx, inferred.
Arguments
a(A): A batch or a scalar.b(B): A batch or a scalar.x(D): A batch or a scalar; at least one argument is a batch.context(_ContextOf[_Triple[A, B, D]]): The family's context; exact values need one.
Returns
Batch[...]: A batch of the broadcast shape.
Raises
Error: When the shapes do not broadcast, or as the scalar function does.
betaln¶
Source: apn_mojo/batch/functions.mojo
def betaln[
A: _MapArgument & ImplicitlyCopyable & Deinitable, B: _MapArgument &
ImplicitlyCopyable & Deinitable, //,
](
a: A,
b: B,
*,
context: _ContextOf[_Pair[A, B]] = _ContextOf[_Pair[A, B]](),
) raises -> Batch[_Pair[A, B]]
log|beta(a, b)| of each pair, like scipy.special.betaln.
Parameters
A(inferred): The type ofa, inferred.B(inferred): The type ofb, inferred.
Arguments
a(A): A batch or a scalar.b(B): A batch or a scalar; at least one argument is a batch.context(_ContextOf[_Pair[A, B]]): The family's context; exact values need one.
Returns
Batch[...]: A batch of the broadcast shape.
Raises
Error: When the shapes do not broadcast, or as the scalar function does.
ceil¶
Source: apn_mojo/batch/functions.mojo
def ceil[T: ImplicitlyCopyable & Deinitable, //](
values: Batch[T],
) raises -> Batch[Integer]
The smallest Integer at least each element, exactly, like numpy.ceil.
Parameters
T(ImplicitlyCopyable & Deinitable, inferred): The element type, inferred.
Arguments
values(Batch[T]): The batch.
Returns
Batch[Integer]: Integers, in a batch of the same shape.
Raises
Error: For a NaN or an infinity among Float elements.
clip¶
Source: apn_mojo/batch/functions.mojo
def clip[
A: _MapArgument & ImplicitlyCopyable & Deinitable, B: _MapArgument &
ImplicitlyCopyable & Deinitable, D: _MapArgument & ImplicitlyCopyable &
Deinitable, //,
](
values: A,
a_min: B,
a_max: D,
) raises -> Batch[_Clipped[A, B, D]]
Each value limited to [a_min, a_max], like numpy.clip.
Exact families stay exact. Real families and balls only.
Parameters
A(inferred): The type ofvalues, inferred.B(inferred): The type ofa_min, inferred.D(inferred): The type ofa_max, inferred.
Arguments
values(A): A batch or a scalar.a_min(B): The lower bounds: a batch or a scalar.a_max(D): The upper bounds: a batch or a scalar; at least one argument is a batch.
Returns
Batch[...]: A batch of the broadcast shape.
Raises
Error: When the shapes do not broadcast, or as the scalar function does.
1 more overload
Each value limited to `[a_min, a_max]`, rounded once with `context`.
def clip[
A: _MapArgument & ImplicitlyCopyable & Deinitable, B: _MapArgument &
ImplicitlyCopyable & Deinitable, D: _MapArgument & ImplicitlyCopyable &
Deinitable, C: ImplicitlyCopyable, //,
](
values: A,
a_min: B,
a_max: D,
*,
context: C,
) raises -> Batch[_WithContext[_Clipped[A, B, D], C]]comb¶
Source: apn_mojo/batch/functions.mojo
def comb[
A: _MapArgument & ImplicitlyCopyable & Deinitable, B: _MapArgument &
ImplicitlyCopyable & Deinitable, //,
](
N: A,
k: B,
*,
repetition: Bool = False,
) raises -> Batch[Integer]
The number of ways to choose k of N things, exactly, like scipy.special.comb with exact=True; 0 when k > N, N < 0 or k < 0.
Parameters
A(inferred): The type ofN, inferred.B(inferred): The type ofk, inferred.
Arguments
N(A): The numbers of things: Integers, a batch or a scalar.k(B): The numbers taken: Integers, a batch or a scalar; at least one argument is a batch.repetition(Bool): Whether a thing may be taken more than once.
Returns
Batch[Integer]: Integers, in a batch of the broadcast shape.
Raises
Error: When the shapes do not broadcast, or a count exceeds the addressable size.
concatenate¶
Source: apn_mojo/batch/joining.mojo
def concatenate[T: ImplicitlyCopyable & Deinitable](
batches: List[Batch[T]],
*,
axis: Int = 0,
) raises -> Batch[T]
The batches joined along an existing axis, like numpy.concatenate.
Parameters
T(ImplicitlyCopyable & Deinitable): The element type, inferred.
Arguments
batches(List[Batch[T]]): One or more batches of one rank, with equal dimensions off the axis.axis(Int): The axis to join along; negative axes count from the end.
Returns
Batch[T]: A new batch holding the values in order.
Raises
Error: For no batches, rank-zero batches, or shapes that differ off the axis.
conjugate¶
Source: apn_mojo/batch/functions.mojo
def conjugate[T: ImplicitlyCopyable & Deinitable, //](
values: Batch[T],
) raises -> Batch[T]
The complex conjugate of every element, like numpy.conjugate.
Parameters
T(ImplicitlyCopyable & Deinitable, inferred): The element type, inferred.
Arguments
values(Batch[T]): The batch.
Returns
Batch[T]: A batch of the same shape.
Raises
Error: Only on a checked size error.
cos¶
Source: apn_mojo/batch/functions.mojo
def cos[T: ImplicitlyCopyable & Deinitable, //](
values: Batch[T],
*,
context: _ContextOf[T] = _ContextOf[T](),
) raises -> Batch[_Mapped[T]]
The cosine of every element, like numpy.cos.
Parameters
T(ImplicitlyCopyable & Deinitable, inferred): The element type, inferred.
Arguments
values(Batch[T]): The batch.context(_ContextOf[T]): The family's context, as for the scalar function; exact elements need one.
Returns
Batch[...]: A batch of the same shape.
Raises
Error: As the scalar function does, with the failing element's position.
cosh¶
Source: apn_mojo/batch/functions.mojo
def cosh[T: ImplicitlyCopyable & Deinitable, //](
values: Batch[T],
*,
context: _ContextOf[T] = _ContextOf[T](),
) raises -> Batch[_Mapped[T]]
The hyperbolic cosine of every element, like numpy.cosh.
Parameters
T(ImplicitlyCopyable & Deinitable, inferred): The element type, inferred.
Arguments
values(Batch[T]): The batch.context(_ContextOf[T]): The family's context, as for the scalar function; exact elements need one.
Returns
Batch[...]: A batch of the same shape.
Raises
Error: As the scalar function does, with the failing element's position.
cumprod¶
Source: apn_mojo/batch/reductions.mojo
def cumprod[T: ImplicitlyCopyable & Deinitable, //](
values: Batch[T],
*,
axis: Optional[Int] = None,
context: _FoldContext[T] = _FoldContext[T](),
) raises -> Batch[T]
Running products along an axis, like numpy.cumprod; without an axis, over the batch flattened in row-major order.
Integer and Rational products are exact. Each Float or Complex product is exact and rounded once, so its working size grows with the number of factors; ball products enclose the exact ones.
Parameters
T(ImplicitlyCopyable & Deinitable, inferred): The element type, inferred.
Arguments
values(Batch[T]): The batch.axis(Optional[Int]): The axis; negative axes count from the end. None flattens first.context(_FoldContext[T]): The rounding of Float, Complex and ball products.
Returns
Batch[T]: A batch of the same shape, or a vector without an axis.
Raises
Error: On an invalid axis, or a context for exact elements.
cumsum¶
Source: apn_mojo/batch/reductions.mojo
def cumsum[T: ImplicitlyCopyable & Deinitable, //](
values: Batch[T],
*,
axis: Optional[Int] = None,
context: _FoldContext[T] = _FoldContext[T](),
) raises -> Batch[T]
Running sums along an axis, like numpy.cumsum; without an axis, over the batch flattened in row-major order.
Integer and Rational sums are exact. Each Float or Complex sum is the exact
sum of its elements rounded once, with context when given, so the result
does not depend on how the elements round; ball sums enclose the exact ones.
Parameters
T(ImplicitlyCopyable & Deinitable, inferred): The element type, inferred.
Arguments
values(Batch[T]): The batch.axis(Optional[Int]): The axis; negative axes count from the end. None flattens first.context(_FoldContext[T]): The rounding of Float, Complex and ball sums.
Returns
Batch[T]: A batch of the same shape, or a vector without an axis.
Raises
Error: On an invalid axis, or a context for exact elements.
digamma¶
Source: apn_mojo/batch/functions.mojo
def digamma[T: ImplicitlyCopyable & Deinitable, //](
values: Batch[T],
*,
context: _ContextOf[T] = _ContextOf[T](),
) raises -> Batch[_Mapped[T]]
The digamma function of every real element or ball, like scipy.special.digamma.
Parameters
T(ImplicitlyCopyable & Deinitable, inferred): The element type, inferred.
Arguments
values(Batch[T]): The batch.context(_ContextOf[T]): The family's context, as for the scalar function; exact elements need one.
Returns
Batch[...]: A batch of the same shape.
Raises
Error: As the scalar function does, with the failing element's position.
divide¶
Source: apn_mojo/batch/functions.mojo
def divide[
A: _MapArgument & ImplicitlyCopyable & Deinitable, B: _MapArgument &
ImplicitlyCopyable & Deinitable, //,
](
a: A,
b: B,
) raises -> Batch[_Arithmetic[A, B, Rational]]
a / b elementwise, like numpy.divide: Integers give exact Rationals; other families as for add.
Parameters
A(inferred): The type ofa, inferred.B(inferred): The type ofb, inferred.
Arguments
a(A): A batch or a scalar.b(B): A batch or a scalar; at least one argument is a batch.
Returns
Batch[...]: A batch of the broadcast shape.
Raises
Error: On division by zero of exact values, when the shapes do not broadcast,
or as the scalar function does.
1 more overload
`a / b` elementwise, each quotient rounded once with `context`, as for `add`.
def divide[
A: _MapArgument & ImplicitlyCopyable & Deinitable, B: _MapArgument &
ImplicitlyCopyable & Deinitable, C: ImplicitlyCopyable, //,
](
a: A,
b: B,
*,
context: C,
) raises -> Batch[_WithContext[_Arithmetic[A, B, Integer], C]]dot¶
Source: apn_mojo/batch/reductions.mojo
def dot[ A: ImplicitlyCopyable, B: ImplicitlyCopyable ](
a: A,
b: B,
) raises -> _DotType[A, B].Value
The sum of pairwise products, exactly.
Elements pair by logical position. Exact families give an exact result; with a Float or Complex operand the exact sum of exact products is rounded once. Integer and Rational operands may pair with Float or Complex ones.
Parameters
A(ImplicitlyCopyable): The first batch type.B(ImplicitlyCopyable): The second batch type.
Arguments
a(A): The first batch.b(B): The second batch, of the same length.
Returns
The dot product; two empty batches give zero.
Raises
Error: When the lengths differ.
1 more overload
The dot product rounded once with `context`.
def dot[ A: ImplicitlyCopyable, B: ImplicitlyCopyable, C: ImplicitlyCopyable ](
a: A,
b: B,
*,
context: C,
) raises -> _DotType[A, B].Valueerf¶
Source: apn_mojo/batch/functions.mojo
def erf[T: ImplicitlyCopyable & Deinitable, //](
values: Batch[T],
*,
context: _ContextOf[T] = _ContextOf[T](),
) raises -> Batch[_Mapped[T]]
The error function of every real element or ball, like scipy.special.erf.
Parameters
T(ImplicitlyCopyable & Deinitable, inferred): The element type, inferred.
Arguments
values(Batch[T]): The batch.context(_ContextOf[T]): The family's context, as for the scalar function; exact elements need one.
Returns
Batch[...]: A batch of the same shape.
Raises
Error: As the scalar function does, with the failing element's position.
erfc¶
Source: apn_mojo/batch/functions.mojo
def erfc[T: ImplicitlyCopyable & Deinitable, //](
values: Batch[T],
*,
context: _ContextOf[T] = _ContextOf[T](),
) raises -> Batch[_Mapped[T]]
The complementary error function of every real element or ball, like scipy.special.erfc.
Parameters
T(ImplicitlyCopyable & Deinitable, inferred): The element type, inferred.
Arguments
values(Batch[T]): The batch.context(_ContextOf[T]): The family's context, as for the scalar function; exact elements need one.
Returns
Batch[...]: A batch of the same shape.
Raises
Error: As the scalar function does, with the failing element's position.
erfi¶
Source: apn_mojo/batch/functions.mojo
def erfi[T: ImplicitlyCopyable & Deinitable, //](
values: Batch[T],
*,
context: _ContextOf[T] = _ContextOf[T](),
) raises -> Batch[_Mapped[T]]
The imaginary error function of every real element or ball, like scipy.special.erfi.
Parameters
T(ImplicitlyCopyable & Deinitable, inferred): The element type, inferred.
Arguments
values(Batch[T]): The batch.context(_ContextOf[T]): The family's context, as for the scalar function; exact elements need one.
Returns
Batch[...]: A batch of the same shape.
Raises
Error: As the scalar function does, with the failing element's position.
erfinv¶
Source: apn_mojo/batch/functions.mojo
def erfinv[T: ImplicitlyCopyable & Deinitable, //](
values: Batch[T],
*,
context: _ContextOf[T] = _ContextOf[T](),
) raises -> Batch[_Mapped[T]]
The inverse error function of every real element or ball, like scipy.special.erfinv.
Parameters
T(ImplicitlyCopyable & Deinitable, inferred): The element type, inferred.
Arguments
values(Batch[T]): The batch.context(_ContextOf[T]): The family's context, as for the scalar function; exact elements need one.
Returns
Batch[...]: A batch of the same shape.
Raises
Error: As the scalar function does, with the failing element's position.
exp¶
Source: apn_mojo/batch/functions.mojo
def exp[T: ImplicitlyCopyable & Deinitable, //](
values: Batch[T],
*,
context: _ContextOf[T] = _ContextOf[T](),
) raises -> Batch[_Mapped[T]]
apn_mojo.exp of every element, like numpy.exp.
Integer, Rational and Float elements give correctly rounded Floats; Complex elements give Complex numbers; balls give balls enclosing the function's values. The other one-argument functions follow the same rules.
Parameters
T(ImplicitlyCopyable & Deinitable, inferred): The element type, inferred.
Arguments
values(Batch[T]): The batch.context(_ContextOf[T]): The family's context, as for the scalar function; exact elements need one.
Returns
Batch[...]: A batch of the same shape.
Raises
Error: As the scalar function does, with the failing element's position.
exp2¶
Source: apn_mojo/batch/functions.mojo
def exp2[T: ImplicitlyCopyable & Deinitable, //](
values: Batch[T],
*,
context: _ContextOf[T] = _ContextOf[T](),
) raises -> Batch[_Mapped[T]]
2**x of every real element or ball, like numpy.exp2.
Parameters
T(ImplicitlyCopyable & Deinitable, inferred): The element type, inferred.
Arguments
values(Batch[T]): The batch.context(_ContextOf[T]): The family's context, as for the scalar function; exact elements need one.
Returns
Batch[...]: A batch of the same shape.
Raises
Error: As the scalar function does, with the failing element's position.
expi¶
Source: apn_mojo/batch/functions.mojo
def expi[T: ImplicitlyCopyable & Deinitable, //](
values: Batch[T],
*,
context: _ContextOf[T] = _ContextOf[T](),
) raises -> Batch[_Mapped[T]]
The exponential integral Ei of every real element or ball, like scipy.special.expi.
Parameters
T(ImplicitlyCopyable & Deinitable, inferred): The element type, inferred.
Arguments
values(Batch[T]): The batch.context(_ContextOf[T]): The family's context, as for the scalar function; exact elements need one.
Returns
Batch[...]: A batch of the same shape.
Raises
Error: As the scalar function does, with the failing element's position.
expm1¶
Source: apn_mojo/batch/functions.mojo
def expm1[T: ImplicitlyCopyable & Deinitable, //](
values: Batch[T],
*,
context: _ContextOf[T] = _ContextOf[T](),
) raises -> Batch[_Mapped[T]]
exp(x) - 1 of every real element or ball, like numpy.expm1.
Parameters
T(ImplicitlyCopyable & Deinitable, inferred): The element type, inferred.
Arguments
values(Batch[T]): The batch.context(_ContextOf[T]): The family's context, as for the scalar function; exact elements need one.
Returns
Batch[...]: A batch of the same shape.
Raises
Error: As the scalar function does, with the failing element's position.
floor¶
Source: apn_mojo/batch/functions.mojo
def floor[T: ImplicitlyCopyable & Deinitable, //](
values: Batch[T],
) raises -> Batch[Integer]
The largest Integer at most each element, exactly, like numpy.floor but with Integer results.
Parameters
T(ImplicitlyCopyable & Deinitable, inferred): The element type, inferred: Integer, Rational or Float.
Arguments
values(Batch[T]): The batch.
Returns
Batch[Integer]: Integers, in a batch of the same shape.
Raises
Error: For a NaN or an infinity among Float elements.
fresnel¶
Source: apn_mojo/batch/functions.mojo
def fresnel[T: ImplicitlyCopyable & Deinitable, //](
values: Batch[T],
*,
context: _ContextOf[T] = _ContextOf[T](),
) raises -> Tuple[Batch[_Mapped[T]], Batch[_Mapped[T]]]
The Fresnel integrals (S, C) of every real element or ball, like scipy.special.fresnel.
Parameters
T(ImplicitlyCopyable & Deinitable, inferred): The element type, inferred.
Arguments
values(Batch[T]): The batch.context(_ContextOf[T]): The family's context, as for the scalar function; exact elements need one.
Returns
Tuple[...]: A batch of each result, both of the same shape.
Raises
Error: As the scalar function does, with the failing element's position.
full¶
Source: apn_mojo/batch/creation.mojo
def full[T: ImplicitlyCopyable & Deinitable](
shape: _Dimensions,
value: T,
) raises -> Batch[T] where conforms_to(T, _BatchElement)
A batch holding value everywhere, like numpy.full; the family and format are the value's.
Parameters
T(ImplicitlyCopyable & Deinitable): The element type, inferred fromvalue.
Arguments
shape(_Dimensions): A length, or a list of dimensions.value(T): The element.
Returns
Batch[T]: A batch of the shape.
Raises
Error: On a negative dimension.
1 more overload
A batch of Integers holding an integer literal of any width.
def full(shape: _Dimensions, value: IntLiteral) raises -> Batch[Integer]gamma¶
Source: apn_mojo/batch/functions.mojo
def gamma[T: ImplicitlyCopyable & Deinitable, //](
values: Batch[T],
*,
context: _ContextOf[T] = _ContextOf[T](),
) raises -> Batch[_Mapped[T]]
The gamma function of every real element or ball, like scipy.special.gamma.
Parameters
T(ImplicitlyCopyable & Deinitable, inferred): The element type, inferred.
Arguments
values(Batch[T]): The batch.context(_ContextOf[T]): The family's context, as for the scalar function; exact elements need one.
Returns
Batch[...]: A batch of the same shape.
Raises
Error: As the scalar function does, with the failing element's position.
gammainc¶
Source: apn_mojo/batch/functions.mojo
def gammainc[
A: _MapArgument & ImplicitlyCopyable & Deinitable, B: _MapArgument &
ImplicitlyCopyable & Deinitable, //,
](
a: A,
x: B,
*,
context: _ContextOf[_Pair[A, B]] = _ContextOf[_Pair[A, B]](),
) raises -> Batch[_Pair[A, B]]
The regularized lower incomplete gamma function, like scipy.special.gammainc.
Parameters
A(inferred): The type ofa, inferred.B(inferred): The type ofx, inferred.
Arguments
a(A): A batch or a scalar.x(B): A batch or a scalar; at least one argument is a batch.context(_ContextOf[_Pair[A, B]]): The family's context; exact values need one.
Returns
Batch[...]: A batch of the broadcast shape.
Raises
Error: When the shapes do not broadcast, or as the scalar function does.
gammaincc¶
Source: apn_mojo/batch/functions.mojo
def gammaincc[
A: _MapArgument & ImplicitlyCopyable & Deinitable, B: _MapArgument &
ImplicitlyCopyable & Deinitable, //,
](
a: A,
x: B,
*,
context: _ContextOf[_Pair[A, B]] = _ContextOf[_Pair[A, B]](),
) raises -> Batch[_Pair[A, B]]
The regularized upper incomplete gamma function, like scipy.special.gammaincc.
Parameters
A(inferred): The type ofa, inferred.B(inferred): The type ofx, inferred.
Arguments
a(A): A batch or a scalar.x(B): A batch or a scalar; at least one argument is a batch.context(_ContextOf[_Pair[A, B]]): The family's context; exact values need one.
Returns
Batch[...]: A batch of the broadcast shape.
Raises
Error: When the shapes do not broadcast, or as the scalar function does.
gammaln¶
Source: apn_mojo/batch/functions.mojo
def gammaln[T: ImplicitlyCopyable & Deinitable, //](
values: Batch[T],
*,
context: _ContextOf[T] = _ContextOf[T](),
) raises -> Batch[_Mapped[T]]
log|gamma(x)| of every real element or ball, like scipy.special.gammaln.
Parameters
T(ImplicitlyCopyable & Deinitable, inferred): The element type, inferred.
Arguments
values(Batch[T]): The batch.context(_ContextOf[T]): The family's context, as for the scalar function; exact elements need one.
Returns
Batch[...]: A batch of the same shape.
Raises
Error: As the scalar function does, with the failing element's position.
hyp1f1¶
Source: apn_mojo/batch/functions.mojo
def hyp1f1[
A: _MapArgument & ImplicitlyCopyable & Deinitable, B: _MapArgument &
ImplicitlyCopyable & Deinitable, D: _MapArgument & ImplicitlyCopyable &
Deinitable, //,
](
a: A,
b: B,
x: D,
*,
context: _ContextOf[_Triple[A, B, D]] = _ContextOf[_Triple[A, B, D]](),
) raises -> Batch[_Triple[A, B, D]]
Kummer's confluent hypergeometric function, like scipy.special.hyp1f1.
Parameters
A(inferred): The type ofa, inferred.B(inferred): The type ofb, inferred.D(inferred): The type ofx, inferred.
Arguments
a(A): A batch or a scalar.b(B): A batch or a scalar.x(D): A batch or a scalar; at least one argument is a batch.context(_ContextOf[_Triple[A, B, D]]): The family's context; exact values need one.
Returns
Batch[...]: A batch of the broadcast shape.
Raises
Error: When the shapes do not broadcast, or as the scalar function does.
hyp2f1¶
Source: apn_mojo/batch/functions.mojo
def hyp2f1[
A: _MapArgument & ImplicitlyCopyable & Deinitable, B: _MapArgument &
ImplicitlyCopyable & Deinitable, D: _MapArgument & ImplicitlyCopyable &
Deinitable, E: _MapArgument & ImplicitlyCopyable & Deinitable, //,
](
a: A,
b: B,
c: D,
x: E,
*,
context: _ContextOf[_Quadruple[A, B, D, E]] = _ContextOf[_Quadruple[A, B, D, E]](),
) raises -> Batch[_Quadruple[A, B, D, E]]
The Gauss hypergeometric function, like scipy.special.hyp2f1.
Parameters
A(inferred): The type ofa, inferred.B(inferred): The type ofb, inferred.D(inferred): The type ofc, inferred.E(inferred): The type ofx, inferred.
Arguments
a(A): A batch or a scalar.b(B): A batch or a scalar.c(D): A batch or a scalar.x(E): A batch or a scalar; at least one argument is a batch.context(_ContextOf[_Quadruple[A, B, D, E]]): The family's context; exact values need one.
Returns
Batch[...]: A batch of the broadcast shape.
Raises
Error: When the shapes do not broadcast, or as the scalar function does.
imag¶
Source: apn_mojo/batch/functions.mojo
def imag[T: ImplicitlyCopyable & Deinitable, //](
values: Batch[T],
) raises -> Batch[_Magnitude[T]]
The imaginary part of every element, like numpy.imag: Floats or balls.
Parameters
T(ImplicitlyCopyable & Deinitable, inferred): The element type, inferred.
Arguments
values(Batch[T]): The batch.
Returns
Batch[...]: A batch of the same shape.
Raises
Error: Only on a checked size error.
lambertw¶
Source: apn_mojo/batch/functions.mojo
def lambertw[T: ImplicitlyCopyable & Deinitable, //](
values: Batch[T],
*,
k: Int = 0,
context: _ContextOf[T] = _ContextOf[T](),
) raises -> Batch[_Mapped[T]]
Branch k of the Lambert W function of every real element or ball, like scipy.special.lambertw.
Parameters
T(ImplicitlyCopyable & Deinitable, inferred): The element type, inferred.
Arguments
values(Batch[T]): The batch.k(Int): The branch.context(_ContextOf[T]): The family's context, as for the scalar function; exact elements need one.
Returns
Batch[...]: A batch of the same shape.
Raises
Error: As the scalar function does, with the failing element's position.
linspace¶
Source: apn_mojo/batch/creation.mojo
def linspace[T: ImplicitlyCopyable & Deinitable](
start: Rational,
stop: Rational,
num: Int = 50,
*,
endpoint: Bool = True,
context: _FormatContext[T] = _FormatContext[T](),
) raises -> Batch[T]
num evenly spaced values from start to stop, like numpy.linspace.
Each element is computed exactly and the last is stop itself, as for
arange.
Parameters
T(ImplicitlyCopyable & Deinitable): The element type: Integer, Rational or Float.
Arguments
start(Rational): The first value.stop(Rational): The last value, or the bound whenendpointis False.num(Int): The number of values.endpoint(Bool): Whetherstopis included.context(_FormatContext[T]): The Float format; 128 bits by default.
Returns
Batch[T]: A vector of num elements.
Raises
Error: For a negative count, or a value that is not whole for Integer elements.
log¶
Source: apn_mojo/batch/functions.mojo
def log[T: ImplicitlyCopyable & Deinitable, //](
values: Batch[T],
*,
context: _ContextOf[T] = _ContextOf[T](),
) raises -> Batch[_Mapped[T]]
The natural logarithm of every element, like numpy.log; Complex elements take the principal branch.
Parameters
T(ImplicitlyCopyable & Deinitable, inferred): The element type, inferred.
Arguments
values(Batch[T]): The batch.context(_ContextOf[T]): The family's context, as for the scalar function; exact elements need one.
Returns
Batch[...]: A batch of the same shape.
Raises
Error: As the scalar function does, with the failing element's position.
log10¶
Source: apn_mojo/batch/functions.mojo
def log10[T: ImplicitlyCopyable & Deinitable, //](
values: Batch[T],
*,
context: _ContextOf[T] = _ContextOf[T](),
) raises -> Batch[_Mapped[T]]
The base-10 logarithm of every real element or ball, like numpy.log10.
Parameters
T(ImplicitlyCopyable & Deinitable, inferred): The element type, inferred.
Arguments
values(Batch[T]): The batch.context(_ContextOf[T]): The family's context, as for the scalar function; exact elements need one.
Returns
Batch[...]: A batch of the same shape.
Raises
Error: As the scalar function does, with the failing element's position.
log1p¶
Source: apn_mojo/batch/functions.mojo
def log1p[T: ImplicitlyCopyable & Deinitable, //](
values: Batch[T],
*,
context: _ContextOf[T] = _ContextOf[T](),
) raises -> Batch[_Mapped[T]]
log(1 + x) of every real element or ball, like numpy.log1p.
Parameters
T(ImplicitlyCopyable & Deinitable, inferred): The element type, inferred.
Arguments
values(Batch[T]): The batch.context(_ContextOf[T]): The family's context, as for the scalar function; exact elements need one.
Returns
Batch[...]: A batch of the same shape.
Raises
Error: As the scalar function does, with the failing element's position.
log2¶
Source: apn_mojo/batch/functions.mojo
def log2[T: ImplicitlyCopyable & Deinitable, //](
values: Batch[T],
*,
context: _ContextOf[T] = _ContextOf[T](),
) raises -> Batch[_Mapped[T]]
The base-2 logarithm of every real element or ball, like numpy.log2.
Parameters
T(ImplicitlyCopyable & Deinitable, inferred): The element type, inferred.
Arguments
values(Batch[T]): The batch.context(_ContextOf[T]): The family's context, as for the scalar function; exact elements need one.
Returns
Batch[...]: A batch of the same shape.
Raises
Error: As the scalar function does, with the failing element's position.
log_ndtr¶
Source: apn_mojo/batch/functions.mojo
def log_ndtr[T: ImplicitlyCopyable & Deinitable, //](
values: Batch[T],
*,
context: _ContextOf[T] = _ContextOf[T](),
) raises -> Batch[_Mapped[T]]
The logarithm of ndtr of every real element or ball, like scipy.special.log_ndtr.
Parameters
T(ImplicitlyCopyable & Deinitable, inferred): The element type, inferred.
Arguments
values(Batch[T]): The batch.context(_ContextOf[T]): The family's context, as for the scalar function; exact elements need one.
Returns
Batch[...]: A batch of the same shape.
Raises
Error: As the scalar function does, with the failing element's position.
max¶
Source: apn_mojo/batch/reductions.mojo
def max[T: ImplicitlyCopyable](values: T) raises -> _ExtremeType[T].Value
The greatest element (numpy's max).
Integers and Rationals compare exactly; a Float batch gives the first greatest element, or its first NaN; a Ball batch gives the ball of the greatest value over all points of its balls.
Parameters
T(ImplicitlyCopyable): The batch type.
Arguments
values(T): A batch, or any selection of one.
Returns
The greatest value.
Raises
Error: When the batch is empty; check len(values) first.
1 more overload
The greatest elements along an axis or a list of axes.
def max[T: ImplicitlyCopyable](
values: T,
*,
var axis: _ReduceAxes,
keepdims: Bool = False,
) raises -> Batch[_ExtremeType[T].Value] where conforms_to(T, _BatchShape)maximum¶
Source: apn_mojo/batch/functions.mojo
def maximum[
A: _MapArgument & ImplicitlyCopyable & Deinitable, B: _MapArgument &
ImplicitlyCopyable & Deinitable, //,
](
a: A,
b: B,
) raises -> Batch[_Arithmetic[A, B, Integer]]
The larger of each pair, like numpy.maximum; a NaN wins, as in numpy.
Exact families stay exact. Real families and balls only.
Parameters
A(inferred): The type ofa, inferred.B(inferred): The type ofb, inferred.
Arguments
a(A): A batch or a scalar.b(B): A batch or a scalar; at least one argument is a batch.
Returns
Batch[...]: A batch of the broadcast shape.
Raises
Error: When the shapes do not broadcast, or as the scalar function does.
1 more overload
The larger of each pair, rounded once with `context`.
def maximum[
A: _MapArgument & ImplicitlyCopyable & Deinitable, B: _MapArgument &
ImplicitlyCopyable & Deinitable, C: ImplicitlyCopyable, //,
](
a: A,
b: B,
*,
context: C,
) raises -> Batch[_WithContext[_Arithmetic[A, B, Integer], C]]min¶
Source: apn_mojo/batch/reductions.mojo
def min[T: ImplicitlyCopyable](values: T) raises -> _ExtremeType[T].Value
The least element (numpy's min).
Integers and Rationals compare exactly; a Float batch gives the first least element, or its first NaN; a Ball batch gives the ball of the least value over all points of its balls.
Parameters
T(ImplicitlyCopyable): The batch type.
Arguments
values(T): A batch, or any selection of one.
Returns
The least value.
Raises
Error: When the batch is empty; check len(values) first.
1 more overload
The least elements along an axis or a list of axes.
def min[T: ImplicitlyCopyable](
values: T,
*,
var axis: _ReduceAxes,
keepdims: Bool = False,
) raises -> Batch[_ExtremeType[T].Value] where conforms_to(T, _BatchShape)minimum¶
Source: apn_mojo/batch/functions.mojo
def minimum[
A: _MapArgument & ImplicitlyCopyable & Deinitable, B: _MapArgument &
ImplicitlyCopyable & Deinitable, //,
](
a: A,
b: B,
) raises -> Batch[_Arithmetic[A, B, Integer]]
The smaller of each pair, like numpy.minimum; as for maximum.
Parameters
A(inferred): The type ofa, inferred.B(inferred): The type ofb, inferred.
Arguments
a(A): A batch or a scalar.b(B): A batch or a scalar; at least one argument is a batch.
Returns
Batch[...]: A batch of the broadcast shape.
Raises
Error: When the shapes do not broadcast, or as the scalar function does.
1 more overload
The smaller of each pair, rounded once with `context`.
def minimum[
A: _MapArgument & ImplicitlyCopyable & Deinitable, B: _MapArgument &
ImplicitlyCopyable & Deinitable, C: ImplicitlyCopyable, //,
](
a: A,
b: B,
*,
context: C,
) raises -> Batch[_WithContext[_Arithmetic[A, B, Integer], C]]multiply¶
Source: apn_mojo/batch/functions.mojo
def multiply[
A: _MapArgument & ImplicitlyCopyable & Deinitable, B: _MapArgument &
ImplicitlyCopyable & Deinitable, //,
](
a: A,
b: B,
) raises -> Batch[_Arithmetic[A, B, Integer]]
a * b elementwise, like numpy.multiply; families as for add.
Parameters
A(inferred): The type ofa, inferred.B(inferred): The type ofb, inferred.
Arguments
a(A): A batch or a scalar.b(B): A batch or a scalar; at least one argument is a batch.
Returns
Batch[...]: A batch of the broadcast shape.
Raises
Error: When the shapes do not broadcast, or as the scalar function does.
1 more overload
`a * b` elementwise, each product rounded once with `context`, as for `add`.
def multiply[
A: _MapArgument & ImplicitlyCopyable & Deinitable, B: _MapArgument &
ImplicitlyCopyable & Deinitable, C: ImplicitlyCopyable, //,
](
a: A,
b: B,
*,
context: C,
) raises -> Batch[_WithContext[_Arithmetic[A, B, Integer], C]]ndtr¶
Source: apn_mojo/batch/functions.mojo
def ndtr[T: ImplicitlyCopyable & Deinitable, //](
values: Batch[T],
*,
context: _ContextOf[T] = _ContextOf[T](),
) raises -> Batch[_Mapped[T]]
The standard normal distribution function of every real element or ball, like scipy.special.ndtr.
Parameters
T(ImplicitlyCopyable & Deinitable, inferred): The element type, inferred.
Arguments
values(Batch[T]): The batch.context(_ContextOf[T]): The family's context, as for the scalar function; exact elements need one.
Returns
Batch[...]: A batch of the same shape.
Raises
Error: As the scalar function does, with the failing element's position.
ndtri¶
Source: apn_mojo/batch/functions.mojo
def ndtri[T: ImplicitlyCopyable & Deinitable, //](
values: Batch[T],
*,
context: _ContextOf[T] = _ContextOf[T](),
) raises -> Batch[_Mapped[T]]
The inverse of ndtr of every real element or ball, like scipy.special.ndtri.
Parameters
T(ImplicitlyCopyable & Deinitable, inferred): The element type, inferred.
Arguments
values(Batch[T]): The batch.context(_ContextOf[T]): The family's context, as for the scalar function; exact elements need one.
Returns
Batch[...]: A batch of the same shape.
Raises
Error: As the scalar function does, with the failing element's position.
ones¶
Source: apn_mojo/batch/creation.mojo
def ones[T: ImplicitlyCopyable & Deinitable](
shape: _Dimensions,
*,
context: _FormatContext[T] = _FormatContext[T](),
) raises -> Batch[T]
A batch of ones, like numpy.ones.
Parameters
T(ImplicitlyCopyable & Deinitable): The element type: Integer, Rational, Float, Complex, Ball or ComplexBall.
Arguments
shape(_Dimensions): A length, or a list of dimensions.context(_FormatContext[T]): The Float or Complex format; 128 bits by default.
Returns
Batch[T]: A batch of the shape, every element one.
Raises
Error: On a negative dimension or an invalid context.
poch¶
Source: apn_mojo/batch/functions.mojo
def poch[
A: _MapArgument & ImplicitlyCopyable & Deinitable, B: _MapArgument &
ImplicitlyCopyable & Deinitable, //,
](
z: A,
m: B,
*,
context: _ContextOf[_Pair[A, B]] = _ContextOf[_Pair[A, B]](),
) raises -> Batch[_Pair[A, B]]
The Pochhammer symbol (z)_m of each pair, like scipy.special.poch.
Parameters
A(inferred): The type ofz, inferred.B(inferred): The type ofm, inferred.
Arguments
z(A): A batch or a scalar.m(B): A batch or a scalar; at least one argument is a batch.context(_ContextOf[_Pair[A, B]]): The family's context; exact values need one.
Returns
Batch[...]: A batch of the broadcast shape.
Raises
Error: When the shapes do not broadcast, or as the scalar function does.
polygamma¶
Source: apn_mojo/batch/functions.mojo
def polygamma[
A: _MapArgument & ImplicitlyCopyable & Deinitable, B: _MapArgument &
ImplicitlyCopyable & Deinitable, //,
](
n: A,
x: B,
*,
context: _ContextOf[_Pair[A, B]] = _ContextOf[_Pair[A, B]](),
) raises -> Batch[_Pair[A, B]]
The nth derivative of digamma at each x, like scipy.special.polygamma; n is Integers.
Parameters
A(inferred): The type ofn, inferred.B(inferred): The type ofx, inferred.
Arguments
n(A): The orders, Integers: a batch or a scalar.x(B): A batch or a scalar; at least one argument is a batch.context(_ContextOf[_Pair[A, B]]): The family's context; exact values need one.
Returns
Batch[...]: A batch of the broadcast shape.
Raises
Error: When the shapes do not broadcast, or as the scalar function does.
pow¶
Source: apn_mojo/batch/functions.mojo
def pow[
A: _MapArgument & ImplicitlyCopyable & Deinitable, B: _MapArgument &
ImplicitlyCopyable & Deinitable, //,
](
base: A,
exponent: B,
*,
context: _ContextOf[_Pair[A, B]] = _ContextOf[_Pair[A, B]](),
) raises -> Batch[_Pair[A, B]]
base ** exponent elementwise, like numpy.power, for real, Complex and ball values.
Parameters
A(inferred): The type ofbase, inferred.B(inferred): The type ofexponent, inferred.
Arguments
base(A): A batch or a scalar.exponent(B): A batch or a scalar; at least one argument is a batch.context(_ContextOf[_Pair[A, B]]): The family's context; exact values need one.
Returns
Batch[...]: A batch of the broadcast shape.
Raises
Error: When the shapes do not broadcast, or as the scalar function does.
pow_int¶
Source: apn_mojo/batch/functions.mojo
def pow_int[
A: _MapArgument & ImplicitlyCopyable & Deinitable, B: _MapArgument &
ImplicitlyCopyable & Deinitable, //,
](
values: A,
exponent: B,
*,
context: _ContextOf[_Pair[A, B]] = _ContextOf[_Pair[A, B]](),
) raises -> Batch[_Pair[A, B]]
Each value to an Integer power, rounded once.
Exact values give Floats under a context; Complex values, Complex numbers; balls, balls.
Parameters
A(inferred): The type ofvalues, inferred.B(inferred): The type ofexponent, inferred.
Arguments
values(A): A batch or a scalar of any family.exponent(B): Integers: a batch or a scalar.context(_ContextOf[_Pair[A, B]]): The family's context; exact values need one.
Returns
Batch[...]: A batch of the broadcast shape.
Raises
Error: When the shapes do not broadcast, or as the scalar function does.
prod¶
Source: apn_mojo/batch/reductions.mojo
def prod[T: ImplicitlyCopyable](values: T) raises -> _ReductionType[T].Value
The exact product of every element.
Parameters
T(ImplicitlyCopyable): The batch type, of Integers or Rationals.
Arguments
values(T): A batch, or any selection of one.
Returns
The product; an empty product is one.
Raises
Error: Only on a checked size error.
1 more overload
Exact products along an axis or a list of axes.
def prod[T: ImplicitlyCopyable](
values: T,
*,
var axis: _ReduceAxes,
keepdims: Bool = False,
) raises -> Batch[_ReductionType[T].Value] where conforms_to(T, _BatchShape)real¶
Source: apn_mojo/batch/functions.mojo
def real[T: ImplicitlyCopyable & Deinitable, //](
values: Batch[T],
) raises -> Batch[_Magnitude[T]]
The real part of every element, like numpy.real: Floats or balls.
Parameters
T(ImplicitlyCopyable & Deinitable, inferred): The element type, inferred.
Arguments
values(Batch[T]): The batch.
Returns
Batch[...]: A batch of the same shape.
Raises
Error: Only on a checked size error.
reciprocal¶
Source: apn_mojo/batch/functions.mojo
def reciprocal[T: ImplicitlyCopyable & Deinitable, //](
values: Batch[T],
) raises -> Batch[ Rational if T == Integer else T ]
1 / x elementwise, like numpy.reciprocal: Integers give exact Rationals, Rationals stay exact, the other families keep theirs.
Parameters
T(ImplicitlyCopyable & Deinitable, inferred): The element type, inferred.
Arguments
values(Batch[T]): The batch.
Returns
Batch[ Rational if T == Integer else T ]: A batch of the same shape.
Raises
Error: As the scalar function does, with the failing element's position.
1 more overload
`1 / x` elementwise, each rounded once with `context`, as for `add`.
def reciprocal[T: ImplicitlyCopyable & Deinitable, C: ImplicitlyCopyable, //](
values: Batch[T],
*,
context: C,
) raises -> Batch[_WithContext[T, C]]rootn¶
Source: apn_mojo/batch/functions.mojo
def rootn[T: ImplicitlyCopyable & Deinitable, //](
values: Batch[T],
n: Int,
*,
context: _ContextOf[T] = _ContextOf[T](),
) raises -> Batch[_Mapped[T]]
The real nth root of every real element or ball.
Parameters
T(ImplicitlyCopyable & Deinitable, inferred): The element type, inferred.
Arguments
values(Batch[T]): The batch.n(Int): The degree of the root.context(_ContextOf[T]): The family's context, as for the scalar function; exact elements need one.
Returns
Batch[...]: A batch of the same shape.
Raises
Error: As the scalar function does, with the failing element's position.
round¶
Source: apn_mojo/batch/functions.mojo
def round[T: ImplicitlyCopyable & Deinitable, //](
values: Batch[T],
) raises -> Batch[Integer]
Each element rounded to the nearest Integer, ties to even, like numpy.round.
Parameters
T(ImplicitlyCopyable & Deinitable, inferred): The element type, inferred.
Arguments
values(Batch[T]): The batch.
Returns
Batch[Integer]: Integers, in a batch of the same shape.
Raises
Error: For a NaN or an infinity among Float elements.
shichi¶
Source: apn_mojo/batch/functions.mojo
def shichi[T: ImplicitlyCopyable & Deinitable, //](
values: Batch[T],
*,
context: _ContextOf[T] = _ContextOf[T](),
) raises -> Tuple[Batch[_Mapped[T]], Batch[_Mapped[T]]]
The hyperbolic sine and cosine integrals of every real element or ball, like scipy.special.shichi.
Parameters
T(ImplicitlyCopyable & Deinitable, inferred): The element type, inferred.
Arguments
values(Batch[T]): The batch.context(_ContextOf[T]): The family's context, as for the scalar function; exact elements need one.
Returns
Tuple[...]: A batch of each result, both of the same shape.
Raises
Error: As the scalar function does, with the failing element's position.
sici¶
Source: apn_mojo/batch/functions.mojo
def sici[T: ImplicitlyCopyable & Deinitable, //](
values: Batch[T],
*,
context: _ContextOf[T] = _ContextOf[T](),
) raises -> Tuple[Batch[_Mapped[T]], Batch[_Mapped[T]]]
The sine and cosine integrals of every real element or ball, like scipy.special.sici.
Parameters
T(ImplicitlyCopyable & Deinitable, inferred): The element type, inferred.
Arguments
values(Batch[T]): The batch.context(_ContextOf[T]): The family's context, as for the scalar function; exact elements need one.
Returns
Tuple[...]: A batch of each result, both of the same shape.
Raises
Error: As the scalar function does, with the failing element's position.
sin¶
Source: apn_mojo/batch/functions.mojo
def sin[T: ImplicitlyCopyable & Deinitable, //](
values: Batch[T],
*,
context: _ContextOf[T] = _ContextOf[T](),
) raises -> Batch[_Mapped[T]]
The sine of every element, like numpy.sin.
Parameters
T(ImplicitlyCopyable & Deinitable, inferred): The element type, inferred.
Arguments
values(Batch[T]): The batch.context(_ContextOf[T]): The family's context, as for the scalar function; exact elements need one.
Returns
Batch[...]: A batch of the same shape.
Raises
Error: As the scalar function does, with the failing element's position.
sin_cos¶
Source: apn_mojo/batch/functions.mojo
def sin_cos[T: ImplicitlyCopyable & Deinitable, //](
values: Batch[T],
*,
context: _ContextOf[T] = _ContextOf[T](),
) raises -> Tuple[Batch[_Mapped[T]], Batch[_Mapped[T]]]
The sines and the cosines of every real element or ball, computed together.
Parameters
T(ImplicitlyCopyable & Deinitable, inferred): The element type, inferred.
Arguments
values(Batch[T]): The batch.context(_ContextOf[T]): The family's context, as for the scalar function; exact elements need one.
Returns
Tuple[...]: A batch of each result, both of the same shape.
Raises
Error: As the scalar function does, with the failing element's position.
sinh¶
Source: apn_mojo/batch/functions.mojo
def sinh[T: ImplicitlyCopyable & Deinitable, //](
values: Batch[T],
*,
context: _ContextOf[T] = _ContextOf[T](),
) raises -> Batch[_Mapped[T]]
The hyperbolic sine of every element, like numpy.sinh.
Parameters
T(ImplicitlyCopyable & Deinitable, inferred): The element type, inferred.
Arguments
values(Batch[T]): The batch.context(_ContextOf[T]): The family's context, as for the scalar function; exact elements need one.
Returns
Batch[...]: A batch of the same shape.
Raises
Error: As the scalar function does, with the failing element's position.
sqrt¶
Source: apn_mojo/batch/functions.mojo
def sqrt[T: ImplicitlyCopyable & Deinitable, //](
values: Batch[T],
*,
context: _ContextOf[T] = _ContextOf[T](),
) raises -> Batch[_Mapped[T]]
The square root of every element, like numpy.sqrt.
Integer, Rational and Float elements give correctly rounded Floats; Complex elements give principal roots; balls give enclosing balls.
Parameters
T(ImplicitlyCopyable & Deinitable, inferred): The element type, inferred.
Arguments
values(Batch[T]): The batch.context(_ContextOf[T]): The family's context, as for the scalar function; exact elements need one.
Returns
Batch[...]: A batch of the same shape.
Raises
Error: As the scalar function does, with the failing element's position.
stack¶
Source: apn_mojo/batch/joining.mojo
def stack[T: ImplicitlyCopyable & Deinitable](
batches: List[Batch[T]],
*,
axis: Int = 0,
) raises -> Batch[T]
The batches joined along a new axis, like numpy.stack.
Parameters
T(ImplicitlyCopyable & Deinitable): The element type, inferred.
Arguments
batches(List[Batch[T]]): One or more batches of one shape.axis(Int): The position of the new axis in the result; negative axes count from the end.
Returns
Batch[T]: A new batch with one more axis.
Raises
Error: For no batches or different shapes.
subtract¶
Source: apn_mojo/batch/functions.mojo
def subtract[
A: _MapArgument & ImplicitlyCopyable & Deinitable, B: _MapArgument &
ImplicitlyCopyable & Deinitable, //,
](
a: A,
b: B,
) raises -> Batch[_Arithmetic[A, B, Integer]]
a - b elementwise, like numpy.subtract; families as for add.
Parameters
A(inferred): The type ofa, inferred.B(inferred): The type ofb, inferred.
Arguments
a(A): A batch or a scalar.b(B): A batch or a scalar; at least one argument is a batch.
Returns
Batch[...]: A batch of the broadcast shape.
Raises
Error: When the shapes do not broadcast, or as the scalar function does.
1 more overload
`a - b` elementwise, each difference rounded once with `context`, as for `add`.
def subtract[
A: _MapArgument & ImplicitlyCopyable & Deinitable, B: _MapArgument &
ImplicitlyCopyable & Deinitable, C: ImplicitlyCopyable, //,
](
a: A,
b: B,
*,
context: C,
) raises -> Batch[_WithContext[_Arithmetic[A, B, Integer], C]]sum¶
Source: apn_mojo/batch/reductions.mojo
def sum[T: ImplicitlyCopyable](values: T) raises -> _SumType[T].Value
Sum every element exactly.
Integer and Rational sums are exact. Float and Complex sums add the stored values exactly and round once, per component for Complex, so the result does not depend on the order of the elements. An empty sum is zero.
Parameters
T(ImplicitlyCopyable): The batch type.
Arguments
values(T): A batch, or any selection of one.
Returns
The total, in the element family.
Raises
Error: When Float formats have different exponent bounds without a context, or on
a checked size error.
2 more overloads
Sum a Float or Complex batch exactly and round once with `context`.
def sum[ T: ImplicitlyCopyable, C: ImplicitlyCopyable ](
values: T,
*,
context: C,
) raises -> _SumType[T].ValueSum exactly along an axis or a list of axes.
def sum[T: ImplicitlyCopyable](
values: T,
*,
var axis: _ReduceAxes,
keepdims: Bool = False,
) raises -> Batch[_SumType[T].Value] where conforms_to(T, _BatchShape)tan¶
Source: apn_mojo/batch/functions.mojo
def tan[T: ImplicitlyCopyable & Deinitable, //](
values: Batch[T],
*,
context: _ContextOf[T] = _ContextOf[T](),
) raises -> Batch[_Mapped[T]]
The tangent of every element, like numpy.tan.
Parameters
T(ImplicitlyCopyable & Deinitable, inferred): The element type, inferred.
Arguments
values(Batch[T]): The batch.context(_ContextOf[T]): The family's context, as for the scalar function; exact elements need one.
Returns
Batch[...]: A batch of the same shape.
Raises
Error: As the scalar function does, with the failing element's position.
tanh¶
Source: apn_mojo/batch/functions.mojo
def tanh[T: ImplicitlyCopyable & Deinitable, //](
values: Batch[T],
*,
context: _ContextOf[T] = _ContextOf[T](),
) raises -> Batch[_Mapped[T]]
The hyperbolic tangent of every element, like numpy.tanh.
Parameters
T(ImplicitlyCopyable & Deinitable, inferred): The element type, inferred.
Arguments
values(Batch[T]): The batch.context(_ContextOf[T]): The family's context, as for the scalar function; exact elements need one.
Returns
Batch[...]: A batch of the same shape.
Raises
Error: As the scalar function does, with the failing element's position.
trunc¶
Source: apn_mojo/batch/functions.mojo
def trunc[T: ImplicitlyCopyable & Deinitable, //](
values: Batch[T],
) raises -> Batch[Integer]
Each element rounded toward zero to an Integer, like numpy.trunc.
Parameters
T(ImplicitlyCopyable & Deinitable, inferred): The element type, inferred.
Arguments
values(Batch[T]): The batch.
Returns
Batch[Integer]: Integers, in a batch of the same shape.
Raises
Error: For a NaN or an infinity among Float elements.
vdot¶
Source: apn_mojo/batch/reductions.mojo
def vdot[ A: ImplicitlyCopyable, B: ImplicitlyCopyable ](
a: A,
b: B,
) raises -> _DotType[A, B].Value
The dot product with the first operand conjugated.
For real operands it equals dot.
Parameters
A(ImplicitlyCopyable): The first batch type.B(ImplicitlyCopyable): The second batch type.
Arguments
a(A): The batch to conjugate.b(B): The second batch, of the same length.
Returns
The sum of conj(a[i]) * b[i], rounded once per component.
Raises
Error: When the lengths differ.
1 more overload
The conjugated dot product rounded once with `context`.
def vdot[ A: ImplicitlyCopyable, B: ImplicitlyCopyable, C: ImplicitlyCopyable ](
a: A,
b: B,
*,
context: C,
) raises -> _DotType[A, B].Valuevmap¶
Source: apn_mojo/batch/mapping.mojo
def vmap[
T: ImplicitlyCopyable & Deinitable, R: ImplicitlyCopyable & Deinitable, //,
function: def(T) raises thin -> R,
](
*__disambiguate: NoneType,
var in_axes: _AxisLevels = {},
var out_axes: _AxisLevels = {},
axis_size: Optional[Int] = None,
var out_shape: _Shapes = {},
) -> _Map[T, R, function]
Map a scalar function over batches.
vmap[f]() returns a mapped function: var mapped = vmap[f](); mapped(xs).
Each argument is a scalar or a batch. A mapped argument passes one slice along
its mapped axis to each call, and a scalar is shared by every call; numeric
results stack into a batch, Boolean results into a Mask, and batch results
into a batch with one more axis. A function of one or two arguments that
returns an Optional number gives two results: the values, with
0 where there is none, and a Mask of where there is one. Name a family's single declaration, such as
vmap[apn_mojo.float.add]; the package-level add is an overload set. Mapped
lengths must agree; a length-one batch is a sequence, not a scalar. Long
mappings with number or Bool results, tuples of them included, run on the
worker pool with the same results and errors.
Limitations
The function cannot capture local variables; pass shared values as
arguments or context=. Functions take one to four positional
arguments, optionally with a keyword-only context, or one flat tuple.
Parameters
T(ImplicitlyCopyable & Deinitable, inferred): The function's argument type, inferred.R(ImplicitlyCopyable & Deinitable, inferred): The function's result type, inferred.function(def(T) raises thin -> R): The scalar function: a nameddefor a lambda that captures nothing.
Arguments
in_axes(var _AxisLevels): The mapped axis of each argument, orNoneto share it; a list sets one entry per nesting level, innermost first.out_axes(var _AxisLevels): Where the new axis goes in each result.axis_size(Optional[Int]): The mapped length when no argument is mapped.out_shape(var _Shapes): The per-call result shape, for empty mappings of batch results.
Returns
The mapped function. .vmap(...) on it adds an outer layer.
14 more overloads
Map a unary function that returns a tuple; each leaf collects separately.
def vmap[
T: ImplicitlyCopyable & Deinitable, R: type_of(Tuple), //, function: def(T)
raises thin -> R,
](
*,
var in_axes: _AxisLevels = {},
var out_axes: _AxisLevels = {},
axis_size: Optional[Int] = None,
var out_shape: _Shapes = {},
) -> _TupleResultMap[T, R, function] where (_MappingLeaf[T])Map a unary function that returns an Optional number: the values (0 where None) and a Mask of where they exist.
def vmap[
T: ImplicitlyCopyable & Deinitable, X: ImplicitlyCopyable & Deinitable, //,
function: def(T) raises thin -> Optional[X],
](
*,
var in_axes: _AxisLevels = {},
var out_axes: _AxisLevels = {},
axis_size: Optional[Int] = None,
var out_shape: _Shapes = {},
) -> _TupleResultMap[T, Tuple[X, Bool], _present[T, X, function]] where (_MappingLeaf[T] and _Present[X])Map a function of one flat tuple, with one `in_axes` entry per leaf.
def vmap[
T: _TupleArgument & ImplicitlyCopyable & Deinitable, R: ImplicitlyCopyable &
Deinitable, //, function: def(T) raises thin -> R, *Ts: ImplicitlyCopyable &
Deinitable,
](
*,
var in_axes: _AxisLevels = {},
var out_axes: _AxisLevels = {},
axis_size: Optional[Int] = None,
var out_shape: _Shapes = {},
) -> _TupleArgumentMap[T, R, function, Tuple[*Ts]] where (T == Tuple[*Ts] and _MappingLeaf[R])Map a function of one flat tuple that returns a tuple.
def vmap[
T: _TupleArgument & ImplicitlyCopyable & Deinitable, R: type_of(Tuple), //,
function: def(T) raises thin -> R, *Ts: ImplicitlyCopyable & Deinitable,
](
*,
var in_axes: _AxisLevels = {},
var out_axes: _AxisLevels = {},
axis_size: Optional[Int] = None,
var out_shape: _Shapes = {},
) -> _TupleArgumentResultMap[T, R, function, Tuple[*Ts]] where (T == Tuple[*Ts])Map a binary function.
def vmap[
T: ImplicitlyCopyable & Deinitable, U: ImplicitlyCopyable & Deinitable, R:
ImplicitlyCopyable & Deinitable, //, function: def(T, U) raises thin -> R,
](
*__disambiguate: NoneType,
var in_axes: _AxisLevels = {},
var out_axes: _AxisLevels = {},
axis_size: Optional[Int] = None,
var out_shape: _Shapes = {},
) -> _Map2[T, U, R, function]Map a binary function that returns a tuple.
def vmap[
T: ImplicitlyCopyable & Deinitable, U: ImplicitlyCopyable & Deinitable, R:
type_of(Tuple), //, function: def(T, U) raises thin -> R,
](
*,
var in_axes: _AxisLevels = {},
var out_axes: _AxisLevels = {},
axis_size: Optional[Int] = None,
var out_shape: _Shapes = {},
) -> _TupleResultMap2[T, U, R, function]Map a binary function that returns an Optional number: the values (0 where None) and a Mask of where they exist.
def vmap[
T: ImplicitlyCopyable & Deinitable, U: ImplicitlyCopyable & Deinitable, X:
ImplicitlyCopyable & Deinitable, //, function: def(T, U) raises thin ->
Optional[X],
](
*,
var in_axes: _AxisLevels = {},
var out_axes: _AxisLevels = {},
axis_size: Optional[Int] = None,
var out_shape: _Shapes = {},
) -> _TupleResultMap2[T, U, Tuple[X, Bool], _present2[T, U, X, function]] where (_Present[X])Map a unary function with a keyword-only `context`, passed to every call.
def vmap[
T: ImplicitlyCopyable & Deinitable, C: ImplicitlyCopyable & Deinitable &
Defaultable, R: ImplicitlyCopyable & Deinitable, //, function: def(T, /, *,
context: C) raises thin -> R,
](
*__disambiguate: NoneType,
var in_axes: _AxisLevels = {},
var out_axes: _AxisLevels = {},
axis_size: Optional[Int] = None,
var out_shape: _Shapes = {},
) -> _MapC[T, C, R, function] where (_Context[C])Map a unary function with a keyword-only `context` that returns a tuple, such as `sici`; each leaf collects separately.
def vmap[
T: ImplicitlyCopyable & Deinitable, C: ImplicitlyCopyable & Deinitable &
Defaultable, R: type_of(Tuple), //, function: def(T, /, *, context: C)
raises thin -> R,
](
*,
var in_axes: _AxisLevels = {},
var out_axes: _AxisLevels = {},
axis_size: Optional[Int] = None,
var out_shape: _Shapes = {},
) -> _TupleResultMapC[T, C, R, function] where (_Context[C])Map a binary function with a keyword-only `context`, passed to every call.
def vmap[
T: ImplicitlyCopyable & Deinitable, U: ImplicitlyCopyable & Deinitable, C:
ImplicitlyCopyable & Deinitable & Defaultable, R: ImplicitlyCopyable &
Deinitable, //, function: def(T, U, /, *, context: C) raises thin -> R,
](
*,
var in_axes: _AxisLevels = {},
var out_axes: _AxisLevels = {},
axis_size: Optional[Int] = None,
var out_shape: _Shapes = {},
) -> _Map2C[T, U, C, R, function] where (_Context[C])Map a ternary function, such as `fma`.
def vmap[
T: ImplicitlyCopyable & Deinitable, U: ImplicitlyCopyable & Deinitable, W:
ImplicitlyCopyable & Deinitable, R: ImplicitlyCopyable & Deinitable, //,
function: def(T, U, W, /) raises thin -> R,
](
*__disambiguate: NoneType,
var in_axes: _AxisLevels = {},
var out_axes: _AxisLevels = {},
axis_size: Optional[Int] = None,
var out_shape: _Shapes = {},
) -> _Map3[T, U, W, R, function]Map a ternary function with a keyword-only `context`, passed to every call.
def vmap[
T: ImplicitlyCopyable & Deinitable, U: ImplicitlyCopyable & Deinitable, W:
ImplicitlyCopyable & Deinitable, C: ImplicitlyCopyable & Deinitable &
Defaultable, R: ImplicitlyCopyable & Deinitable, //, function: def(T, U, W,
/, *, context: C) raises thin -> R,
](
*,
var in_axes: _AxisLevels = {},
var out_axes: _AxisLevels = {},
axis_size: Optional[Int] = None,
var out_shape: _Shapes = {},
) -> _Map3C[T, U, W, C, R, function] where (_Context[C])Map a four-argument function.
def vmap[
T: ImplicitlyCopyable & Deinitable, U: ImplicitlyCopyable & Deinitable, W:
ImplicitlyCopyable & Deinitable, X: ImplicitlyCopyable & Deinitable, R:
ImplicitlyCopyable & Deinitable, //, function: def(T, U, W, X, /) raises
thin -> R,
](
*__disambiguate: NoneType,
var in_axes: _AxisLevels = {},
var out_axes: _AxisLevels = {},
axis_size: Optional[Int] = None,
var out_shape: _Shapes = {},
) -> _Map4[T, U, W, X, R, function]Map a four-argument function, such as hyp2f1, with a keyword-only `context`, passed to every call.
def vmap[
T: ImplicitlyCopyable & Deinitable, U: ImplicitlyCopyable & Deinitable, W:
ImplicitlyCopyable & Deinitable, X: ImplicitlyCopyable & Deinitable, C:
ImplicitlyCopyable & Deinitable & Defaultable, R: ImplicitlyCopyable &
Deinitable, //, function: def(T, U, W, X, /, *, context: C) raises thin ->
R,
](
*,
var in_axes: _AxisLevels = {},
var out_axes: _AxisLevels = {},
axis_size: Optional[Int] = None,
var out_shape: _Shapes = {},
) -> _Map4C[T, U, W, X, C, R, function] where (_Context[C])where¶
Source: apn_mojo/batch/functions.mojo
def where[
A: _MapArgument & ImplicitlyCopyable & Deinitable, B: _MapArgument &
ImplicitlyCopyable & Deinitable, //,
](
condition: Mask,
a: A,
b: B,
) raises -> Batch[_Chosen[A, B]]
Elements of a where condition holds and of b elsewhere, like numpy.where.
a and b broadcast to the mask's shape; masks do not broadcast. A
scalar converts exactly to the batch's family, as vmap converts it.
Parameters
A(inferred): The type ofa, inferred.B(inferred): The type ofb, inferred.
Arguments
condition(Mask): The mask.a(A): A batch or a scalar.b(B): A batch or a scalar; at least one ofaandbis a batch.
Returns
Batch[...]: A batch of the mask's shape.
Raises
Error: When a or b does not broadcast to the mask's shape.
zeros¶
Source: apn_mojo/batch/creation.mojo
def zeros[T: ImplicitlyCopyable & Deinitable](
shape: _Dimensions,
*,
context: _FormatContext[T] = _FormatContext[T](),
) raises -> Batch[T]
A batch of zeros, like numpy.zeros.
Parameters
T(ImplicitlyCopyable & Deinitable): The element type: Integer, Rational, Float, Complex, Ball or ComplexBall.
Arguments
shape(_Dimensions): A length, or a list of dimensions.context(_FormatContext[T]): The Float or Complex format; 128 bits by default.
Returns
Batch[T]: A batch of the shape, every element zero.
Raises
Error: On a negative dimension or an invalid context.
zeta¶
Source: apn_mojo/batch/functions.mojo
def zeta[
A: _MapArgument & ImplicitlyCopyable & Deinitable, B: _MapArgument &
ImplicitlyCopyable & Deinitable, //,
](
x: A,
q: B,
*,
context: _ContextOf[_Pair[A, B]] = _ContextOf[_Pair[A, B]](),
) raises -> Batch[_Pair[A, B]]
The Hurwitz zeta function zeta(x, q) of each pair, like scipy.special.zeta.
Parameters
A(inferred): The type ofx, inferred.B(inferred): The type ofq, inferred.
Arguments
x(A): A batch or a scalar.q(B): A batch or a scalar; at least one argument is a batch.context(_ContextOf[_Pair[A, B]]): The family's context; exact values need one.
Returns
Batch[...]: A batch of the broadcast shape.
Raises
Error: When the shapes do not broadcast, or as the scalar function does.
1 more overload
The Riemann zeta function of every real element or ball, like `scipy.special.zeta` without `q`.
def zeta[T: ImplicitlyCopyable & Deinitable, //](
x: Batch[T],
*,
context: _ContextOf[T] = _ContextOf[T](),
) raises -> Batch[_Mapped[T]]Structs¶
Batch¶
Source: apn_mojo/batch/value.mojo
struct Batch[T: ImplicitlyCopyable & Deinitable](
ImplicitlyCopyable,
Iterable,
IterableOwned,
Sized,
Writable,
_ComplexBatchOperand,
_ExactBatchComparison,
_FloatBatchOperand,
_TensorSequence,
_BatchShape,
_MapArgument,
)
A batch of numbers of any rank, with value semantics.
Rank and shape are run-time values, so one type serves rank-zero scalars, vectors, matrices and higher-rank tensors. Slices, axis selections, transposes and reshapes share the values without copying them, and a batch never sees a later update to another one: updates publish new storage.
| Operation | Contract |
|---|---|
a + b, a - b, a * b, a / b |
Elementwise; shapes broadcast from the trailing axis; scalars broadcast |
a // b, a % b, a ** b, bit operations |
Elementwise on equal shapes or a scalar |
==, !=, <, <=, >, >= |
A Mask of the broadcast shape |
values[i], values[i, j] |
One element; negative indices count from the end |
values[a:b:c] |
A vector over the flat row-major values, sharing them |
values[mask] |
A new vector of the selected values |
values[selection] = source |
A scalar fills; a sequence must match the selection's length |
values += x and the other compound forms |
Keep the destination's shape and formats; unchanged on error |
Two vectors combine only at equal lengths: a length-one vector is a sequence,
not a scalar. Exact / returns a Batch[Rational]. A failed operation reports
the earliest failing element and publishes nothing. Long operations run on
the worker pool with the same results and errors.
Named math functions are scalar; apply them with vmap, as in
vmap[apn_mojo.float.sqrt]()(values).
Limitations
A native number cannot be the left operand of a comparison; write
batch > native. Batches are not hash keys. A slice keeps its whole
source alive; copy a small slice with Batch[T](values.to_list()).
Parameters
T(ImplicitlyCopyable & Deinitable): The element type:Integer,Rational,Float,ComplexorBall. ABatch[Ball]is a container, mapped withvmap; it has no operators or JSON.
Implements
Copyable, ImplicitlyCopyable, Iterable, IterableOwned, Sized, Writable, _BatchShape, _ComplexBatchArithmetic, _ComplexBatchOperand, _ExactBatchComparison, _FloatBatchArithmetic, _FloatBatchOperand, _MapArgument, _TensorSequence
Aliases
FloatBatchComplexBatchElementIteratorTypeIteratorOwnedType
Batch.__init__ { #Batch.init .api-name }¶
def __init__(out self, values: List[T], *, shape: List[Int]) raises
A batch with a shape, from row-major values.
Arguments
values(List[T]): The values in row-major order.shape(List[Int]): The dimensions, outermost first; their product must equal the value count.
Raises
Error: When the shape does not match the value count.
def __init__(out self, values: List[T]) raises
A vector of the values; a batch of the same family shares its values.
Arguments
values(List[T]): The values.
Raises
Error: Only on a checked size error.
8 more overloads
A batch with a shape, from a finite iterable consumed once.
def __init__[I: Iterable](
out self,
values: I,
*,
shape: List[Int],
) raises where I != List[T]A batch with a shape, from an owned finite iterable.
def __init__(
out self,
var values: Some[IterableOwned],
*,
shape: List[Int],
) raisesA rank-zero batch holding one value; it fills a selection when assigned.
def __init__(out self, value: T)A rank-zero batch from an integer literal of any width.
def __init__(out self, value: IntLiteral) raisesA rank-zero batch from a value that widens exactly to the element type.
def __init__[V: ImplicitlyCopyable & Deinitable](
out self,
value: V,
) raises where (V != T and _Exact[T, V])An empty vector.
def __init__(out self)A vector from a finite iterable, consumed once; another batch converts each element.
def __init__[I: Iterable](out self, values: I) raisesA vector from an owned finite iterable.
def __init__(out self, var values: Some[IterableOwned]) raisesBatch.at¶
def at(self, axis: Int, index: Int) raises -> Self
Select one index along an axis, dropping that axis.
Arguments
axis(Int): The axis; negative counts from the end.index(Int): The index along it; negative counts from the end.
Returns
Self: A batch of one rank lower, sharing the values; values.at(0, i) is row i.
Raises
Error: When the axis or index is out of range.
Batch.ceil¶
def ceil(self,) raises -> Batch[Integer] where T == Rational or T == Float
Round each element toward positive infinity.
Returns
Batch[Integer]: A Batch[Integer] of the same shape.
Raises
Error: For an infinite or NaN element, naming its index.
Batch.conjugate¶
def conjugate(self) raises -> Self where T == Complex
The complex conjugate of each element.
Returns
Self: A Complex batch of the same shape.
Raises
Error: Only on a checked size error.
Batch.floor¶
def floor(self,) raises -> Batch[Integer] where T == Rational or T == Float
Round each element toward negative infinity.
Returns
Batch[Integer]: A Batch[Integer] of the same shape.
Raises
Error: For an infinite or NaN element, naming its index.
Batch.from_iterable¶
def from_iterable[I: Iterable](values: I, *, shape: List[Int]) raises -> Self
A batch from a finite iterable, consumed once.
Parameters
I(Iterable): The iterable type; its items must convert exactly to the element type.
Arguments
values(I): The values in row-major order.shape(List[Int]): The dimensions, outermost first; their product must equal the value count.
Returns
Self: The batch.
Raises
Error: When an item does not convert or the shape does not match.
4 more overloads
A batch with a shape from an owned finite iterable.
def from_iterable(
var values: Some[IterableOwned],
*,
shape: List[Int],
) raises -> SelfA vector of the values.
def from_iterable(values: List[T]) raises -> SelfA vector from a finite iterable, consumed once.
def from_iterable[I: Iterable](values: I) raises -> SelfA vector from an owned finite iterable.
def from_iterable(var values: Some[IterableOwned]) raises -> SelfBatch.from_json¶
def from_json(
text: String,
*,
limits: Optional[ConversionLimits] = None,
) raises -> Self
Read a batch from its JSON record.
Version 1 holds a vector; version 2 adds an explicit shape.
Arguments
text(String): The JSON record.limits(Optional[ConversionLimits]): Optional per-call conversion limits; seeConversionLimits.
Returns
Self: An independent batch.
Raises
Error: When the text is not exactly that schema, naming the byte offset and element.
Batch.from_native¶
def from_native(values: List[Int], *, shape: List[Int]) raises -> Self
A batch from native integers.
Arguments
values(List[Int]): The values in row-major order.shape(List[Int]): The dimensions, outermost first; their product must equal the value count.
Returns
Self: The batch.
Raises
Error: When the shape does not match the value count.
3 more overloads
A Float or Complex batch with a shape, from native numbers.
def from_native[dtype: DType](
values: List[SIMD[dtype, 1]],
*,
shape: List[Int],
) raises -> Self where (T == Float or T == Complex)A vector from native integers.
def from_native(values: List[Int]) raises -> SelfA Float or Complex vector from native numbers.
def from_native[ dtype: DType ](
values: List[SIMD[dtype, 1]],
) raises -> Self where (T == Float or T == Complex)Batch.imag¶
def imag(self) raises -> Batch[Float] where T == Complex
The imaginary part of each element.
Returns
Batch[Float]: A Batch[Float] of the same shape.
Raises
Error: Only on a checked size error.
Batch.is_finite¶
def is_finite(self,) raises -> Mask where T == Float or T == Complex
Whether each element is finite.
Returns
Mask: A Mask of the same shape.
Raises
Error: Only on a checked size error.
Batch.is_infinite¶
def is_infinite(self,) raises -> Mask where T == Float or T == Complex
Whether each element is infinite.
Returns
Mask: A Mask of the same shape.
Raises
Error: Only on a checked size error.
Batch.is_nan¶
def is_nan(self) raises -> Mask where T == Float or T == Complex
Whether each element is NaN.
Returns
Mask: A Mask of the same shape.
Raises
Error: Only on a checked size error.
Batch.is_zero¶
def is_zero(self) raises -> Mask where T == Float or T == Complex
Whether each element is zero.
Returns
Mask: A Mask of the same shape.
Raises
Error: Only on a checked size error.
Batch.item¶
def item(self) raises -> T
The only element.
Returns
T: The element of a batch of size one.
Raises
Error: When the size is not one.
Batch.magnitude_bit_length¶
def magnitude_bit_length(self) raises -> Self
The bit length of each element's absolute value.
Returns
Self: A batch of the same shape.
Raises
Error: Only on a checked size error.
Batch.ndim¶
def ndim(self) -> Int
The rank.
Returns
Int: The number of axes; zero for a rank-zero batch.
Batch.real¶
def real(self) raises -> Batch[Float] where T == Complex
The real part of each element.
Returns
Batch[Float]: A Batch[Float] of the same shape.
Raises
Error: Only on a checked size error.
Batch.reshape¶
def reshape(self, shape: List[Int]) raises -> Self
The same values with another shape, in row-major order.
Arguments
shape(List[Int]): The new dimensions; the size must not change.
Returns
Self: A batch sharing the values when they form one strided run.
Raises
Error: When the sizes differ.
Batch.shape¶
def shape(self) -> List[Int]
The dimensions.
Returns
List[Int]: A copy of the shape, outermost axis first.
Batch.sign¶
def sign(self) raises -> Batch[Integer]
The sign of each element: -1, 0 or 1.
Returns
Batch[Integer]: A Batch[Integer] of the same shape.
Raises
Error: For a NaN element.
Batch.signbit¶
def signbit(self) raises -> Mask where T == Float
Whether each element is negative, including negative zero.
Returns
Mask: A Mask of the same shape.
Raises
Error: Only on a checked size error.
Batch.size¶
def size(self) -> Int
The number of elements.
Returns
Int: The product of the dimensions; len(values) is the same.
Batch.slice¶
def slice(
self,
axis: Int,
start: Optional[Int] = None,
stop: Optional[Int] = None,
step: Int = 1,
) raises -> Self
Slice one axis, keeping the rank.
Arguments
axis(Int): The axis; negative counts from the end.start(Optional[Int]): The first index, as in Python slicing.stop(Optional[Int]): The end index, as in Python slicing.step(Int): The step; not zero.
Returns
Self: A batch sharing the values.
Raises
Error: When the axis is out of range or the step is zero.
Batch.to_integer_exact¶
def to_integer_exact(self) raises -> Batch[Integer] where T == Float
Convert each integral element to Integer.
Returns
Batch[Integer]: A Batch[Integer] of the same shape.
Raises
Error: For a fractional, infinite or NaN element, naming its index.
Batch.to_json¶
def to_json(self, *, limits: Optional[ConversionLimits] = None) raises -> String
Write the JSON record: version 1 for a vector, version 2 with its shape otherwise.
Selections write their values in logical order.
Arguments
limits(Optional[ConversionLimits]): Optional per-call conversion limits; seeConversionLimits.
Returns
String: Compact canonical JSON.
Raises
Error: When the output exceeds limits.
Batch.to_list¶
def to_list(self) raises -> List[T]
The values, in row-major order.
Returns
List[T]: An independent list.
Raises
Error: Only on a checked size error.
Batch.to_native¶
def to_native[dtype: DType](
self,
*,
rounding: RoundingMode = RoundingMode.nearest_even,
) raises -> List[SIMD[dtype, 1]]
The values as native numbers in row-major order, the counterpart of from_native; shape() gives the dimensions.
To float64, float32, float16 or bfloat16, each value is rounded
once with rounding, subnormals included; beyond the range it becomes an
infinity. A Ball converts through its midpoint. To an integer type, each
value converts exactly: a value that is not whole or does not fit raises.
Parameters
dtype(DType): The native type.
Arguments
rounding(RoundingMode): The rounding mode for floating-point types.
Returns
List[SIMD[dtype, 1]]: One native value per element.
Raises
Error: For an integer type, at the first value that is not a whole number in
its range.
Batch.transpose¶
def transpose(self, axes: List[Int]) raises -> Self
Permute the axes.
Arguments
axes(List[Int]): A permutation of the axes; negative entries count from the end.
Returns
Self: A batch sharing the values.
Raises
Error: When axes is not a permutation.
1 more overload
Reverse the axes, sharing the values.
def transpose(self) raises -> SelfBatch.trunc¶
def trunc(self,) raises -> Batch[Integer] where T == Rational or T == Float
Round each element toward zero.
Returns
Batch[Integer]: A Batch[Integer] of the same shape.
Raises
Error: For an infinite or NaN element, naming its index.
Batch.write_to¶
def write_to(self, mut writer: Some[Writer])
Write the values as nested lists, such as [[1, 2], [3, 4]], in full.
Arguments
writer(mut Some[Writer]): The destination.
Mask¶
Source: apn_mojo/batch/mask.mojo
struct Mask(
ImplicitlyCopyable,
Iterable,
IterableOwned,
Sized,
Writable,
_MapArgument,
)
An immutable batch of Booleans, returned by batch comparisons.
A mask has a shape like a batch. batch[mask] needs a mask of exactly the
batch's shape and selects the true positions in row-major order. &, |, ^
and ~ combine masks of equal shape; nothing broadcasts. A comparison mask keeps
its values when the numbers change.
Limitations
Bool(mask) always raises as ambiguous: ask any() or all(). Use &
and |, not and and or, to combine masks.
Implements
Copyable, ImplicitlyCopyable, Iterable, IterableOwned, Sized, Writable, _MapArgument
Aliases
IteratorTypeIteratorOwnedType
Mask.__init__ { #Mask.init .api-name }¶
def __init__[I: Iterable](out self, values: I) raises
A mask from a finite iterable of Booleans.
Parameters
I(Iterable): The iterable type.
Arguments
values(I): The Booleans.
Raises
Error: When an item is not a Boolean.
4 more overloads
A mask from an owned finite iterable of Booleans.
def __init__(out self, var values: Some[IterableOwned]) raisesAn empty mask.
def __init__(out self)A vector mask from a list of Booleans.
def __init__(out self, values: List[Bool]) raisesA mask with a shape, from row-major Booleans.
def __init__(out self, values: List[Bool], *, shape: List[Int]) raisesMask.all¶
def all(self) -> Bool
Whether every entry is true.
Returns
Bool: True for an empty mask.
Mask.any¶
def any(self) -> Bool
Whether some entry is true.
Returns
Bool: False for an empty mask.
Mask.at¶
def at(self, axis: Int, index: Int) raises -> Self
Select one index along an axis, dropping that axis.
Arguments
axis(Int): The axis.index(Int): The index along it.
Returns
Self: A mask of one rank lower.
Raises
Error: When the axis or index is out of range.
Mask.count¶
def count(self) -> Int
The number of true entries.
Returns
Int: The count; the length of batch[mask].
Mask.from_iterable¶
def from_iterable[I: Iterable](values: I) raises -> Self
A mask from a finite iterable of Booleans.
Parameters
I(Iterable): The iterable type.
Arguments
values(I): The Booleans.
Returns
Self: The mask.
Raises
Error: When an item is not a Boolean.
1 more overload
A mask from an owned finite iterable of Booleans.
def from_iterable(var values: Some[IterableOwned]) raises -> SelfMask.item¶
def item(self) raises -> Bool
The only entry.
Returns
Bool: The entry of a mask of size one.
Raises
Error: When the size is not one.
Mask.ndim¶
def ndim(self) -> Int
The rank.
Returns
Int: The number of axes.
Mask.reshape¶
def reshape(self, shape: List[Int]) raises -> Self
The same entries with another shape.
Arguments
shape(List[Int]): The new dimensions; the size must not change.
Returns
Self: The reshaped mask.
Raises
Error: When the sizes differ.
Mask.shape¶
def shape(self) -> List[Int]
The dimensions.
Returns
List[Int]: A copy of the shape.
Mask.size¶
def size(self) -> Int
The number of entries.
Returns
Int: The product of the dimensions.
Mask.to_list¶
def to_list(self) raises -> List[Bool]
The entries, in row-major order.
Returns
List[Bool]: An independent list.
Raises
Error: Only on a checked size error.
Mask.transpose¶
def transpose(self, axes: List[Int]) raises -> Self
Permute the axes.
Arguments
axes(List[Int]): A permutation of the axes.
Returns
Self: The transposed mask.
Raises
Error: When axes is not a permutation.
1 more overload
Reverse the axes.
def transpose(self) raises -> SelfMask.write_to¶
def write_to(self, mut writer: Some[Writer])
Write the entries as nested lists of True and False.
Arguments
writer(mut Some[Writer]): The destination.
NumPy-style functions¶
Import batch with from apn_mojo import batch for calls familiar from
numpy or jax.numpy: batch.exp(xs), batch.atan2(ys, xs), and
batch.where(mask, a, b). Scalar functions stay at the package root, keeping
apn_mojo.exp(x) and batch.exp(xs) unambiguous.
"""NumPy-style functions of batches: apn_mojo.batch."""
from apn_mojo import ArithmeticContext, Batch, Float, FloatFormat, Integer, batch
def main() raises:
var grid = batch.arange[Integer](6).reshape([2, 3])
print("grid + row:", batch.add(grid, Batch[Integer]([10, 20, 30])))
print("halves:", batch.divide(grid, 2))
print("large kept:", batch.where(grid > 2, grid, 0))
print("running sums:", batch.cumsum(grid, axis=1))
print("stacked shape:", batch.stack([grid, grid]).shape())
var c = ArithmeticContext(format=FloatFormat.binary64())
var xs = batch.linspace[Float](0, 1, 5, context=c)
print("xs:", xs)
print("exp:", batch.exp(xs))
print("atan2(1, xs):", batch.atan2(1, xs))
print("exp(5), rounded once:", batch.exp(grid, context=c)[1, 2])
Run from the repository root pixi run mojo run -I src docs/examples/batch_functions.mojo
Output
grid + row: [[10, 21, 32], [13, 24, 35]]
halves: [[0, 1/2, 1], [3/2, 2, 5/2]]
large kept: [[0, 0, 0], [3, 4, 5]]
running sums: [[0, 1, 3], [3, 7, 12]]
stacked shape: [2, 2, 3]
xs: [0.0, 0.25, 0.5, 0.75, 1.0]
exp: [1.0, 1.2840254166877414, 1.6487212707001282, 2.117000016612675, 2.718281828459045]
atan2(1, xs): [1.5707963267948966, 1.3258176636680326, 1.1071487177940904, 0.9272952180016122, 0.7853981633974483]
exp(5), rounded once: 148.4131591025766
| Group | Functions |
|---|---|
| Arithmetic | add, subtract, multiply, divide, reciprocal, maximum, minimum, clip, abs |
| Powers and roots | sqrt, pow, pow_int, rootn |
| Elementary | exp, expm1, exp2, log, log1p, log2, log10, sin, cos, tan, sin_cos, asin, acos, atan, atan2, sinh, cosh, tanh, asinh, acosh, atanh |
| Special | gamma, gammaln, digamma, polygamma, beta, betaln, poch, erf, erfc, erfi, erfinv, ndtr, log_ndtr, ndtri, expi, sici, shichi, fresnel, lambertw, zeta, gammainc, gammaincc, betainc, hyp1f1, hyp2f1 |
| Integers and parts | floor, ceil, trunc, round, comb, conjugate, real, imag, angle |
| Selection | where |
| Constructors | zeros, ones, full, arange, linspace |
| Joining | concatenate, stack |
| Reductions | sum, prod, min, max, argmin, argmax, dot, vdot, cumsum, cumprod |
Each elementwise function uses vmap to call the same family declarations as
its root counterpart. Every element therefore follows the scalar function's
rules:
- Families. A function takes the families that declare it. Integer and
Rational elements give correctly rounded Floats under a
context=, Complex elements give Complex numbers, and Ball and ComplexBall elements give enclosing balls. The arithmetic functions keep exact families exact without a context, as the operators do:batch.divideof Integer batches gives Rationals. - Broadcasting. Elementwise arguments align trailing dimensions, which
must match or have size one. This includes one-element vectors, unlike
operators on two vectors. A scalar argument, including an integer literal
of any width, is shared by every element; pass a native
IntasInteger(n), as forvmap. Each elementwise function needs at least one batch argument.wherebroadcasts its value arguments to the mask's shape; the mask itself does not broadcast. - Shapes. Results have the broadcast shape; a function with two results,
such as
sin_cos, returns a batch of each. - Errors. A failing element raises the scalar function's error with that element's position.
Choose a constructor's element type with a parameter, much as you would use
dtype in NumPy:
batch.zeros[Integer]([2, 3]), batch.ones[Float](4, context=c).
arange and linspace compute each element exactly; a Float element is the
exact value rounded once, so no rounding error accumulates along a range.
cumsum and cumprod keep each running Float or Complex value exact and round
it once, as sum does.
Shapes and axis views¶
Batch[T] holds Integer, Rational, Float, Complex, Ball, or ComplexBall values.
Rank and shape are runtime properties: one type serves rank-zero scalars, vectors,
matrices, and higher-rank arrays. Ball batches have container operations,
mapping through vmap, binary operations through lift, and the NumPy-style
functions, such as batch.add. Real Ball batches support min and max;
both ball families support cumsum and cumprod. Use lift for other
folds. Ball batches have no arithmetic operators or batch JSON.
"""Shapes and axis views."""
from apn_mojo import Batch, Integer
def main() raises:
var matrix = Batch[Integer]([1, 2, 3, 4, 5, 6], shape=[2, 3])
print(matrix.shape(), matrix.ndim(), len(matrix))
var columns = matrix.transpose()
matrix[0, 1] = 20
print(matrix)
print(columns)
print(matrix.at(1, -1))
print(matrix.slice(1, step=-1))
for i in range(matrix.shape()[0]):
print(matrix.at(0, i))
print(matrix.reshape([6]))
var scalar = Batch[Integer]([42], shape=[])
print(scalar.item())
Run from the repository root pixi run mojo run -I src docs/examples/ranked_batches.mojo
Output
[2, 3] 2 6
[[1, 20, 3], [4, 5, 6]]
[[1, 4], [2, 5], [3, 6]]
[3, 6]
[[3, 20, 1], [6, 5, 4]]
[1, 20, 3]
[4, 5, 6]
[1, 20, 3, 4, 5, 6]
42
Dimensions can have size zero. Printing uses nested brackets for the shape. Slices, axis views, transpositions, and suitable reshapes share element storage. Some operations on noncontiguous views gather elements into a new buffer. All returned batches keep their values when the source is updated.
Higher-rank arithmetic¶
Integer, rational, float, and complex batch operators use the corresponding scalar arithmetic and broadcast compatible shapes.
"""Arithmetic with broadcasting across ranks."""
from apn_mojo import Batch, Integer
def main() raises:
var matrix = Batch[Integer]([1, 2, 3, 4, 5, 6], shape=[2, 3])
var row = Batch[Integer]([10, 20, 30])
print(matrix + row)
print(matrix / 2)
print(matrix > 3)
print((matrix + row)[1, 2])
print(matrix.slice(1, step=-1) * 2)
Run from the repository root pixi run mojo run -I src docs/examples/ranked_arithmetic.mojo
Output
[[11, 22, 33], [14, 25, 36]]
[[1/2, 1, 3/2], [2, 5/2, 3]]
[[False, False, False], [True, True, True]]
36
[[6, 4, 2], [12, 10, 8]]
Broadcasting aligns trailing dimensions, which must match or have size one.
For example, [columns] can broadcast over [rows, columns]. Two rank-one
vectors are a special case: they must have equal lengths, even if one length
is one. A scalar operand broadcasts over every element.
Result families follow scalar rules. Dividing integer batches gives rationals; mixing with floats or complex numbers gives those rounded families. Comparisons return a mask. In-place updates can broadcast the right operand but cannot expand the destination shape.
Masks must match shape and do not broadcast. Selection returns values in
row-major order. Reductions accept axes and optional keepdims; see their
individual declarations.
| Interface | Two vectors of lengths n and 1, where n > 1 |
|---|---|
| Arithmetic operators | Shape error |
Elementwise batch.*, such as batch.add or batch.atan2 |
Broadcast to length n |
| A direct lifted call | Broadcast to length n |
vmap over both vector axes |
Extent mismatch; share a scalar instead |
The vectorization tutorial demonstrates these differences.
Mapping functions¶
Create a mapped function with vmap[f]() and call it with batch or scalar
arguments. The function is a compile-time parameter. Use a family declaration
such as apn_mojo.float.add, rather than the overloaded apn_mojo.add.
"""Mapping a scalar function over batches."""
from apn_mojo import Batch, Integer, vmap
def polynomial(x: Integer) raises -> Integer:
return x * x + 2 * x + 1
def difference(x: Integer, y: Integer) raises -> Integer:
return x - y
def positive(x: Integer) raises -> Bool:
return x > 0
def shifted(values: Batch[Integer]) raises -> Batch[Integer]:
# A function can map over its own batch parameters; scalars are shared.
return vmap[difference]()(values, 3)
def main() raises:
var values = Batch[Integer]([-2, -1, 0, 1, 2])
var evaluate = vmap[polynomial]()
print(evaluate(values))
print(evaluate(values[::-2]))
print(vmap[difference](in_axes=(0, None))(values, Integer(3)))
print(shifted(values))
print(values[vmap[positive]()(values)])
Run from the repository root pixi run mojo run -I src docs/examples/vmap.mojo
Output
[1, 0, 1, 4, 9]
[9, 1, 1]
[-5, -4, -3, -2, -1]
[-5, -4, -3, -2, -1]
[1, 2]
in_axes chooses the mapped axis of each argument, defaulting to zero.
A scalar is shared across calls; in_axes=None shares an entire batch.
Scalars are library numbers, masks, Bools and literals. Convert a native
number first, such as Integer(n) for an Int: if natives were accepted, a
wide integer literal could bind as an Int and lose its value.
Mapped axes must have equal extents. A context= argument is forwarded to
functions that accept it.
Functions can be named or lambdas, must use supported signatures, and must be pure. Captured local closures are not supported. Pass runtime values as arguments instead.
Rows, columns and tensor results¶
"""Mapping over rows and columns, with tensor results."""
from apn_mojo import Batch, Integer, sum_sequential, vmap
def total(row: Batch[Integer]) raises -> Integer:
return sum_sequential(row)
def reverse(row: Batch[Integer]) raises -> Batch[Integer]:
return row[::-1]
def shift(row: Batch[Integer], offset: Integer) raises -> Batch[Integer]:
return row + offset
def main() raises:
var matrix = Batch[Integer].from_iterable([1, 2, 3, 4, 5, 6], shape=[2, 3])
print(vmap[total]()(matrix))
print(vmap[total](in_axes=1)(matrix))
print(vmap[reverse](out_axes=-1)(matrix))
print(vmap[shift](in_axes=(0, None))(matrix, Integer(10)))
var empty = Batch[Integer].from_iterable(List[Int](), shape=[0, 3])
print(vmap[reverse](out_shape=[3])(empty).shape())
Run from the repository root pixi run mojo run -I src docs/examples/vmap_tensors.mojo
Output
[6, 15]
[5, 7, 9]
[[3, 6], [2, 5], [1, 4]]
[[11, 12, 13], [14, 15, 16]]
[0, 3]
A batch parameter receives the dimensions left after mapping removes an axis,
such as a row or column of a matrix. Scalar parameters receive elements.
out_axes positions the new mapped axis, defaulting to zero. Batch results
must all have the same shape. Empty mappings need out_shape when no call can
supply the shape.
Tuple results¶
"""Mapped functions that return tuples."""
from apn_mojo import Batch, Integer, vmap
def square(x: Integer) raises -> Integer:
return x * x
def describe(x: Integer) raises -> Tuple[Integer, Integer, Bool]:
return (x * x, x * x * x, x > 2)
def affine_value(value: Integer, scale: Integer, offset: Integer) raises -> Integer:
return value * scale + offset
def affine(parts: Tuple[Integer, Integer, Integer]) raises -> Integer:
return affine_value(parts[0], parts[1], parts[2])
def main() raises:
var values = Batch[Integer]([1, 2, 3])
var result = vmap[describe]()(values)
print(result[0])
print(result[1])
print(result[2].count())
print(vmap[affine](in_axes=(0, None, None))((values, Integer(2), Integer(10))))
print(vmap[affine]()((values, values, values)))
var matrix = Batch[Integer].from_iterable(range(1, 7), shape=[2, 3])
print(vmap[square]().vmap()(matrix))
print(vmap[affine](in_axes=(0, None, None)).vmap(in_axes=(0, None, None))((matrix, Integer(2), Integer(10))))
print(vmap[describe]().vmap()(matrix)[2])
Run from the repository root pixi run mojo run -I src docs/examples/vmap_structures.mojo
Output
[1, 4, 9]
[1, 8, 27]
1
[12, 14, 16]
[2, 6, 12]
[[1, 4, 9], [16, 25, 36]]
[[12, 14, 16], [18, 20, 22]]
[[False, False, True], [True, True, True]]
Functions can accept or return flat tuples. Each leaf can have its own axis setting. A failure discards all partial fields. Nested tuple trees are not supported by the current mapping interface.
Optional results¶
Supported one- or two-argument functions returning an Optional number produce
a value batch and a mask. Missing values use the family's zero as a placeholder;
select with the mask before treating them as results. For example,
var roots, found = vmap[integer.iroot_exact]()(xs, 2) gives roots and their
presence mask. This Optional form takes no context keyword; use a plain
function with fixed numerical settings when needed.
Composing mappings¶
"""Composing mapped functions."""
from apn_mojo import Batch, Integer, vmap
def combine(x: Integer, y: Integer) raises -> Tuple[Integer, Integer]:
return (x + y, x * y)
def tuple_sum(pair: Tuple[Integer, Integer]) raises -> Integer:
return pair[0] + pair[1]
def main() raises:
var matrix = Batch[Integer].from_iterable(range(1, 7), shape=[2, 3])
var shared = Batch[Integer]([10, 20, 30])
var result = vmap[combine]().vmap(in_axes=(0, None), out_axes=(0, -1))(matrix, shared)
print(result[0])
print(result[1])
print(vmap[tuple_sum]().vmap(in_axes=(0, None))((matrix, shared)))
var empty = matrix.slice(0, stop=0)
var empty_result = vmap[combine]().vmap(out_shape=([3], [3]))(empty, empty)
print(empty_result[0].shape())
print(empty_result[1].shape())
Run from the repository root pixi run mojo run -I src docs/examples/vmap_composition.mojo
Output
[[11, 22, 33], [14, 25, 36]]
[[10, 40], [40, 100], [90, 180]]
[[11, 22, 33], [14, 25, 36]]
[0, 3]
[0, 3]
Nest mappings with .vmap(...) or a list of axes. For example,
vmap[f](in_axes=[0, 1]) corresponds to
vmap[f](in_axes=0).vmap(in_axes=1). Eligible nested mappings use one flat
execution plan; other signatures use the general mapping path.
Parallel mapping¶
"""Parallel mapping of a long batch."""
from apn_mojo import Batch, Float, FloatFormat, ArithmeticContext, vmap
from apn_mojo import float
def hypotenuse(x: Float, y: Float) raises -> Float:
return float.sqrt(x * x + y * y)
def main() raises:
var context = ArithmeticContext(format=FloatFormat(64))
var xs = Batch[Float]([Float(i, context=context) for i in range(1, 2001)])
var ys = Batch[Float]([Float(2 * i, context=context) for i in range(1, 2001)])
# vmap[f]() returns the mapped function; long batches use worker threads.
var lengths_of = vmap[hypotenuse]()
var lengths = lengths_of(xs, ys)
print(lengths[0], lengths[1999])
# Library functions work the same way, keyword-only context included.
var sums = vmap[float.add]()(xs, ys, context=context)
print(sums[0], sums[1999])
# Every element equals the scalar call.
print(hypotenuse(xs[1999], ys[1999]) == lengths[1999])
print(lengths_of(ys, xs)[0] == lengths[0])
Run from the repository root pixi run mojo run -I src docs/examples/vmap_parallel.mojo
Output
2.2360679774997896964 4472.135954999579393
3.0 6000.0
True
True
Long eligible mappings use the worker pool. Number and Bool results, including flat tuples of them, can use this path with numeric or batch arguments, across strided, reversed, transposed, and nested layouts.
Short runs and single-core environments stay on the caller thread. So do
batch-returning functions, functions taking Bools or Masks, tuple arguments
with more than three leaves, and mappings with axis_size or out_shape.
See execution.
Parallel evaluation gives the same results and lowest failing logical index as sequential evaluation. The executor may call a failed element again to obtain its error message, so mapped functions must be pure.
Pinned-compiler limits¶
The interface supports up to three separate positional arguments, or one flat tuple, with the documented context forms. Use a tuple adapter for a wider argument list. A function must accept and return types that mapping supports; a custom result struct is not automatically a new batch element type.
Boolean mappings¶
Bool outputs become shaped masks. Masks support shape inspection, reshaping,
transposition, indexing, .any(), .all(), and .count().
Construction and iteration¶
Construct from a list of family values, or use a family's native-input helpers,
such as Batch[Integer].from_native([1, 2, 3]). Iteration reads a snapshot in
logical order and returns independent element values. Container operations
are shared across all supported families; conversion helpers vary by family.
Printing¶
print(values) and String(values) use each family's text form, nested to
match the shape. Integer vectors print like [4, -1, 0], Float vectors like
[0.25, 1.0]. Empty vectors print as [].
Native values¶
Use to_list() to keep the APN element types, or to_native[dtype]() to
convert to native numbers. Both produce a flat list; preserve shape()
separately for a multidimensional result.
values.to_native[DType.float64]() returns the elements as a
List[Float64] in row-major order, the counterpart of from_native.
To float64, float32, float16 or bfloat16 each
value is rounded once, subnormals included, with rounding= (nearest-even by
default); beyond the range it becomes an infinity. Balls convert through their
midpoints. To an integer type each value converts exactly, and a value that
is not a whole number in range raises; round first with batch.round or
batch.floor. Integer output is available for Integer, Rational, and Float
batches; Ball batches export only floating-point midpoints. Take batch.real
and batch.imag of a Complex batch first.
"""Export row-major native values, keeping shape separately."""
from std.collections import Array
from apn_mojo import Ball, Batch, Complex, Rational, batch
def main() raises:
var values = Batch[Rational]([Rational(1, 3), Rational(2, 3), Rational(3), Rational(4)]).reshape([2, 2])
var shape = values.shape()
var native = values.to_native[DType.float64]()
print("shape:", shape)
print("native values:", native)
# Array's length is a compile-time parameter; check a runtime batch's size.
if len(native) != 4:
raise Error("This export needs exactly four elements.")
var fixed = Array[Float64, 4](fill=0)
for i in range(4):
fixed[i] = native[i]
print("fixed array, first and last:", fixed[0], fixed[3])
print("integer floors:", batch.floor(values).to_native[DType.int64]())
var complex_values = Batch[Complex]([Complex(1, 2), Complex(3, -4)])
print("native real parts:", batch.real(complex_values).to_native[DType.float64]())
print("native imaginary parts:", batch.imag(complex_values).to_native[DType.float64]())
var uncertain = Batch[Ball]([Ball(3, Rational(1, 8))])
print("midpoints only, uncertainty discarded:", uncertain.to_native[DType.float64]())
Run from the repository root pixi run mojo run -I src docs/examples/batch_native.mojo
Output
shape: [2, 2]
native values: [0.3333333333333333, 0.6666666666666666, 3.0, 4.0]
fixed array, first and last: 0.3333333333333333 4.0
integer floors: [0, 0, 3, 4]
native real parts: [1.0, 3.0]
native imaginary parts: [2.0, -4.0]
midpoints only, uncertainty discarded: [3.0]
The returned List has a runtime length. Mojo's Array[T, N] requires N
at compile time, so the example checks the length before filling a fixed-size
array. Converting a Ball's midpoint discards its uncertainty; keep the balls
or export endpoints separately when the receiving calculation needs bounds.
For a ComplexBall batch, extract the real and imaginary Ball batches first.
Threads¶
Long operations and mappings run on a pool of POSIX threads, with the results
and errors of a sequential loop. set_num_threads(n) sets how many threads
take part, the calling thread included, and get_num_threads() reads it;
n may exceed the core count, and 1 runs everything on the calling thread.
The default is the APN_MOJO_NUM_THREADS environment variable when it is a
positive integer, else one thread per usable physical core. The call never
waits: made while an operation runs, even from inside a mapped function, it
takes effect from the next operation.
"""Choose the thread count without changing numerical results."""
from apn_mojo import Integer, batch, get_num_threads, set_num_threads
def main() raises:
set_num_threads(2)
print("configured threads:", get_num_threads())
var values = batch.arange[Integer](1000)
var total = batch.sum(values)
set_num_threads(1)
print("calling thread only:", get_num_threads())
print("same total:", batch.sum(values) == total)
Run from the repository root pixi run mojo run -I src docs/examples/batch_threads.mojo
Output
configured threads: 2
calling thread only: 1
same total: True
To choose the default at launch, set the variable before starting your own program:
APN_MOJO_NUM_THREADS=4 pixi run --locked mojo run -I src your_program.mojo
The environment is read when the pool is first initialized; use
set_num_threads to change the count during a run. The setting is
process-wide. Changes wait for a running operation to finish, and growing
the pool restarts its workers. Configure it outside mapped callbacks.
The count is the available participation limit, not a promise that every
call uses that many threads. Short work, unsupported mapping signatures,
and calls made while the pool is busy can run on the caller. On platforms
without POSIX-thread support, get_num_threads() returns 1. On Linux,
the default also respects the process's CPU affinity.
Operators and shapes¶
Operators depend on the family: integer batches have bitwise operations and
integer division, while rational, float, and complex batches expose their own
arithmetic. Integer / produces a rational batch. Read the generated
declarations for mixed operands and return types. Ball batches use mapped
ball functions instead of operators.
Indexing and assignment¶
Use the source forms documented for the destination family. A scalar fills
the selection; a replacement sequence must match it. For overlapping
assignments, the batch reads the old source values before changing the
destination. Compound selection updates such as values[mask] += 1 follow
the same rule.
Sharing and copies¶
Views share element storage but have their own layout. Updating either batch
leaves the other unchanged. A retained view can keep a large backing buffer
alive. To detach a small integer selection, for example, construct
Batch[Integer](values.to_list()).
Empty inputs and failures¶
An empty mapping makes no scalar calls, so there is no scalar domain check. Shape, axis, and configuration checks can still fail. An element failure leaves update destinations unchanged and reports the earliest failing logical index.
Rational batches¶
Batch[Rational] provides exact fractional arithmetic, comparisons, and
updates. JSON stores canonical fractions, using version 1 for vectors and
version 2 with a shape for other ranks.
Float batches¶
"""Float batches and their formats."""
from apn_mojo import Batch, Float, FloatFormat, ArithmeticContext, Mask
def main() raises:
var context = ArithmeticContext(format=FloatFormat(64))
var values = Batch[Float](
[
Float("1.5", context=context),
Float("-0"),
Float(7),
]
)
print(values)
var saved = values[::-1]
values[1:] = values[:-1]
print(values)
print(saved)
print(values[1].precision())
print(saved[Mask([True, False, True])])
Run from the repository root pixi run mojo run -I src docs/examples/float_batches.mojo
Output
[1.5, -0.0, 7.0]
[1.5, 1.5, -0.0]
[7.0, -0.0, 1.5]
64
[7.0, 1.5]
"""Float batch arithmetic."""
from apn_mojo import Batch, Float, FloatFormat, ArithmeticContext, Rational
from apn_mojo import float, vmap
def main() raises:
var values = Batch[Float]([Float("1.5"), Float("-2.25"), Float("0.5")])
print(values + Float("0.5"))
print(values * Rational(2, 3))
print(values[values > Float(0)])
var context = ArithmeticContext(format=FloatFormat(3))
print(vmap[float.add]()(values[::-1], Float("0.25"), context=context))
Run from the repository root pixi run mojo run -I src docs/examples/float_batch_arithmetic.mojo
Output
[2.0, -1.75, 1.0]
[1.0, -1.5, 0.333333333333333333333333333333333333334]
[1.5, 0.5]
[0.8, -2.0, 1.8]
Each Float element retains its format. Integer and rational operands keep their exact value until the operation's final rounding.
Unary operations and powers¶
"""Float batch unary operations and powers."""
from apn_mojo import (
Batch,
Float,
Integer,
float,
vmap,
ArithmeticContext,
FloatFormat,
)
def main() raises:
var values = Batch[Float](
[Float("1.5"), Float(-2), Float.zero(negative=True), Float.nan()]
)
print(-values)
print(vmap[float.abs]()(values))
print(values[~values.is_nan()].sign())
var bases = Batch[Float]([Float("1.5"), Float(-2), Float("0.5")])
print(bases ** Batch[Integer]([2, 3, -1]))
print(Float(2) ** Batch[Integer]([-2, -1, 0, 1, 2]))
var narrow = ArithmeticContext(format=FloatFormat(3))
print(vmap[float.pow_int]()(bases, -1, context=narrow))
Run from the repository root pixi run mojo run -I src docs/examples/float_batch_powers.mojo
Output
[-1.5, 2.0, 0.0, nan]
[1.5, 2.0, 0.0, nan]
[1, -1, 0]
[2.25, -8.0, 2.0]
[0.25, 0.5, 1.0, 2.0, 4.0]
[0.6, -0.5, 2.0]
Unary operations preserve formats. Powers take a scalar integer exponent or an integer batch and round each element once.
Strict functions and integral conversion¶
"""Strict Float batch functions and integral conversion."""
from apn_mojo import (
Batch,
Float,
Integer,
FloatFormat,
ArithmeticContext,
square,
ldexp,
fma,
float,
vmap,
)
def main() raises:
var values = Batch[Float]([Float("1.5"), Float(2), Float(3)])
print(vmap[float.sqrt]()(vmap[square]()(values)))
print(vmap[ldexp]()(values, Batch[Integer]([-1, 0, 2])))
print(values.floor())
print(values.ceil())
var context = ArithmeticContext(format=FloatFormat(3))
var x = Batch[Float]([Float("1.5")])
print(vmap[fma]()(x, Float("1.5"), -2, context=context))
print(vmap[square]()(x, context=context) - Float(2))
Run from the repository root pixi run mojo run -I src docs/examples/float_batch_functions.mojo
Output
[1.5, 2.0, 3.0]
[0.75, 2.0, 12.0]
[1, 2, 3]
[2, 2, 3]
[0.25]
[0.0]
Use batch.sqrt and the other NumPy-style functions, or map family functions
such as fma with vmap.
Batch .floor(), .ceil(), and .trunc() return integer batches.
Float compound updates¶
"""Float compound updates."""
from apn_mojo import Batch, Float, Integer, ArithmeticContext, FloatFormat
def main() raises:
var values = Batch[Float](
[
Float(1, context=ArithmeticContext(format=FloatFormat(3))),
Float(1, context=ArithmeticContext(format=FloatFormat(8))),
]
)
var saved = values[:]
var increment = Float(
"0.125", context=ArithmeticContext(format=FloatFormat(16))
)
# Each lane rounds to its own destination precision, not the RHS precision.
values += increment
print(values)
print(values[0].precision(), values[1].precision())
print(saved)
values /= 2
values **= Batch[Integer]([-1, -2])
print(values)
Run from the repository root pixi run mojo run -I src docs/examples/float_batch_updates.mojo
Output
[1.0, 1.125]
3 8
[1.0, 1.0]
[2.0, 3.16]
Compound operators keep destination formats and use nearest-even rounding.
Float batch JSON¶
"""Float batch JSON."""
from apn_mojo import (
Batch,
Float,
FloatFormat,
ArithmeticContext,
ConversionLimits,
)
def main() raises:
var small = ArithmeticContext(format=FloatFormat(3))
var wide = ArithmeticContext(format=FloatFormat(80))
var values = Batch[Float](
[
Float("1.5", context=small),
Float.zero(negative=True, context=wide),
Float.infinity(context=small),
]
)
var saved = values[::-1]
var limits = ConversionLimits(
max_values=3,
max_digits=100,
max_input_bytes=4096,
max_output_bytes=4096,
max_allocated_bytes=65536,
)
var text = saved.to_json(limits=limits)
values[0] = Float(99)
var restored = Batch[Float].from_json(text, limits=limits)
print(restored)
print(
restored[0].precision(),
restored[1].precision(),
restored[2].precision(),
)
print(restored.to_json() == text)
print(saved.to_json() == text)
Run from the repository root pixi run mojo run -I src docs/examples/float_batch_json.mojo
Output
[inf, -0.0, 1.5]
3 80 3
True
True
JSON preserves each value's format and the default format for empty results.
Complex batches¶
"""Complex batches and masks."""
from apn_mojo import Batch, Complex, Mask
def main() raises:
var values = Batch[Complex]([Complex(1, 2), Complex(3, 4), Complex(5, 6)])
var saved = values[::-1]
values[:] = saved
print(values[0] == Complex(5, 6))
values[Mask([True, False, True])] = Complex(0, -1)
print(values[0] == Complex(0, -1), values[1] == Complex(3, 4))
print(saved[0] == Complex(5, 6), saved[2] == Complex(1, 2))
var real_inputs = Batch[Complex].from_native([4, -1, 0])
var count = 0
for value in real_inputs:
if value.imag().is_zero():
count += 1
print(count)
print(len(Batch[Complex](saved[::2])))
Run from the repository root pixi run mojo run -I src docs/examples/complex_batches.mojo
Output
True
True True
True True
3
2
"""Complex batch arithmetic."""
from apn_mojo import (
Batch,
Complex,
Integer,
ArithmeticContext,
FloatFormat,
complex,
vmap,
)
def main() raises:
var samples = Batch[Complex]([Complex(1, 2), Complex(3, 4)])
var weights = Batch[Integer].from_native([2, 3])
var weighted = samples * weights
print(weighted[0] == Complex(2, 4), weighted[1] == Complex(9, 12))
var saved = samples[::-1]
var shifted = 2 - saved
print(shifted[0] == Complex(-1, -4))
print((samples != Complex(0)).all(), (samples == samples[:]).all())
var precise = vmap[complex.multiply]()(
samples, saved, context=ArithmeticContext(format=FloatFormat(256))
)
print(precise[0] == Complex(-5, 10), precise[0].real_format().precision())
print(samples[0] == Complex(1, 2), saved[0] == Complex(3, 4))
Run from the repository root pixi run mojo run -I src docs/examples/complex_batch_arithmetic.mojo
Output
True True
True
True True
True 256
True True
"""Complex batch magnitudes and square roots."""
from apn_mojo import Batch, Complex, Integer, Float, norm_sqr
from apn_mojo import complex, vmap
def main() raises:
var samples = Batch[Complex](
[Complex(3, 4), Complex(-4, Float.zero(negative=True))]
)
var saved = samples[::-1]
print(
vmap[norm_sqr]()(samples)[0] == Float(25),
vmap[complex.abs]()(samples)[0] == Float(5),
)
var roots = vmap[complex.sqrt]()
print(roots(samples)[0] == Complex(2, 1), roots(saved)[0] == Complex(0, -2))
print(samples.conjugate()[0] == Complex(3, -4))
print(samples.real()[0] == Float(3), samples.imag()[0] == Float(4))
print(samples.is_finite().all(), samples.is_zero().any())
var exponents = Batch[Integer].from_native([2, -1])
var powers = vmap[complex.pow_int]()(samples, exponents)
print(powers[0] == Complex(-7, 24))
var cycle = Complex(0, 1) ** exponents
print(cycle[0] == Complex(-1, 0), cycle[1] == Complex(0, -1))
print(samples[0] == Complex(3, 4))
Run from the repository root pixi run mojo run -I src docs/examples/complex_batch_functions.mojo
Output
True True
True True
True
True True
True False
True
True True
True
Complex batch operations follow the scalar rules and round each output component once. Multiplication and similar operations use both input parts together to produce each component.
Complex compound updates¶
"""Complex compound updates."""
from apn_mojo import Batch, Complex, Integer
def main() raises:
var values = Batch[Complex]([Complex(3, 4), Complex(1, -2)])
var saved = values[:]
values += Batch[Integer].from_native([1, 2])
values -= 1
values *= Complex(0, 1)
values /= Complex(0, 1)
values **= 2
print(values[0] == Complex(-7, 24))
print(values[1] == Complex(0, -8))
print(saved[0] == Complex(3, 4))
print(saved[1] == Complex(1, -2))
Run from the repository root pixi run mojo run -I src docs/examples/complex_batch_updates.mojo
Output
True
True
True
True
Compound operators keep the destination component formats. The destination changes only after the whole update succeeds.
Complex batch JSON¶
"""Complex batch JSON."""
from apn_mojo import (
Batch,
Complex,
Float,
FloatFormat,
ArithmeticContext,
ComplexContext,
ConversionLimits,
)
def main() raises:
var settings = ComplexContext(
real=ArithmeticContext(format=FloatFormat(65)),
imag=ArithmeticContext(format=FloatFormat(257)),
)
var values = Batch[Complex](
[
Complex("1.5-2j", context=settings),
Complex(Float.zero(negative=True), Float.infinity()),
]
)
var saved = values[::-1]
var limits = ConversionLimits(
max_values=2,
max_digits=1000,
max_input_bytes=4096,
max_output_bytes=4096,
max_allocated_bytes=65536,
)
var text = saved.to_json(limits=limits)
values[0] = Complex(99)
var restored = Batch[Complex].from_json(text, limits=limits)
print(restored)
print(
restored[1].real_format().precision(),
restored[1].imag_format().precision(),
)
print(restored.to_json() == text)
print(saved.to_json() == text)
Run from the repository root pixi run mojo run -I src docs/examples/complex_batch_json.mojo
Output
[Complex(-0.0, inf), Complex(1.5, -2.0)]
65 257
True
True
JSON preserves both component values, their formats, and container defaults.
Function and arithmetic rules¶
Vectors need equal lengths, scalars broadcast, and higher-rank arrays follow the shape rules above. Use family functions for mapping and root functions for reductions.
Masks¶
A Mask stores boolean values and a shape. It can come from a comparison,
a mapped predicate, or boolean input.
Iteration¶
Masks iterate in logical order. They retain their bits when numerical inputs change, rather than reevaluating a predicate.
Boolean operations¶
Combine same-shaped masks with &, |, ^, and ~. A direct conversion to
Bool raises; use .any() or .all() for a condition.
Selection¶
batch[mask] collects matching values. batch[mask] = source updates those
positions while preserving the destination shape.