Module Ymir_units.Quantity
Values in a unit.
A quantity is a value, its payload, in a Unit.t. Its operations are a quantity's algebra. Within one dimension quantities form a vector space: add and sub convert the second argument to the first's unit, and scaling is map. Across dimensions they form a group: mul, div, pow and root compute the unit. times and per change the unit and keep the payload. These six raise Invalid_argument as the unit algebra does (Algebra). A number leaves a quantity through a named unit, with value:
let v = Quantity.v Unit.(kilo metre / second) speeds let v_si = Quantity.value Unit.(metre / second) v (* times 1000 *) let below = Nx.less (Quantity.value Unit.metre a) (Quantity.value Unit.metre b)
Quantity is an Nx.Ptree.S: a quantity reports its unit as one Walk.case holding the unit's canonical text, then walks its payload, both at the quantity's path. jit, vmap, scan and jvp carry units with no rule of their own, and a compiled function compiles once per unit. A scan carry keeps the initial carry's unit: a step that computes in another unit returns convert (unit init) q. A record of quantities walks each field with walk; Nx.Ptree.cast and Nx.Ptree.Payload then act on the payloads and keep the units.
Gradients. A unit fixes the coordinates its payload is in, and every transformation acts on payloads in them. A gradient is a direction in the parameters' own coordinates, the value whose Nx.Ptree.dot with a tangent is the derivative along it: a quantity's gradient is held in its parameter's unit, and its payload is the derivative per that unit. For the derivative per another unit, differentiate at the parameter converted to that unit:
let per_s = Unit.(one / second) let at_si = { x with h0 = Quantity.convert per_s x.h0 } let d_h0 = Quantity.value per_s (Rune.grad p loss at_si).h0
Never convert the gradient with value: it multiplies by the factor, where the derivative divides by it. A vjp cotangent is held in its result's unit, likewise. A structure's units are the scale an optimiser or sampler steps in.
Dtypes. A float payload converts by one rounded factor. A complex payload scales its real and imaginary parts by the real factor, so an infinite part never meets a zero imaginary factor. An integer payload converts only by an integer factor, and raises rather than wrap, int4 and uint4 included. A bool or bit tensor has no unit: v and map refuse one, and value raises on one that Nx.Ptree.cast or Nx.Ptree.Payload built.
Arithmetic on payloads (mul, div, pow, add, sub, map) follows nx: integers wrap and truncate, and float elements overflow and underflow as their dtype does.
type !'p tThe type for values 'p in a unit.
val walk : ('a, 'b, 'd, 'e) Nx.Ptree.Walk.cursor -> 'a t -> 'b twalk c q reports q's unit with one Walk.case holding its canonical text, then walks its payload with Walk.leaf, both at c's path. A quantity's structure is Nx.Ptree.instantiate (module Quantity.Structure).
module Structure : Nx.Ptree.S with type ('a, 'd) t = 'a tStructure is walk as a structure whose arrays are all positions.
Constructors
v u x is x in u. Raises Invalid_argument "Quantity.v: a bool tensor has no unit" if x is a bool tensor, and with a bit tensor for a bit one.
Queries
value u q is q's payload in u: the payload itself when unit q is u, otherwise the payload times Unit.ratio d (unit q) u, with d the payload's dtype. The result is within 1.5 ulps of the exact product, and eager and compiled calls agree bit for bit on one device.
Raises Invalid_argument as Unit.ratio does, its message naming Quantity.value, as in "Quantity.value: the factor from 1e-35 kg s^-2 to kg s^-2 is 1e-35, which is 0 in float16". An integer payload raises when an element times the factor leaves the dtype, through Nx.check: at once on the host, at the next wait on another set, on a concrete payload or a traced one, as in "Quantity.value: element [41] of an int32 payload overflows converting 1e3 m to m (factor 1000)", or "an int32 payload overflows" for a scalar. The index is the first such element's in C order.
Changing the unit
convert u q is v u (value u q). It raises as value does, its message naming Quantity.convert.
Maps
map f q is f of q's payload, in q's unit. It is correct when f (c·x) = c·f x for every positive scalar c, such as a scaling, a sum or a mean; nothing checks f. Raises Invalid_argument "Quantity.map: the function returns a bool tensor" if f does, and with a bit tensor for a bit one.
val map2 :
(('a, 'b, 'd) Nx.t -> ('a, 'b, 'd) Nx.t -> ('a, 'b, 'd) Nx.t) ->
('a, 'b, 'd) Nx.t t ->
('a, 'b, 'd) Nx.t t ->
('a, 'b, 'd) Nx.t tmap2 f a b is f of a's payload and value (unit a) b, in a's unit. It is correct when f (c·x) (c·y) = c·f x y. For an integer payload, b converts only when unit b / unit a is an integer: convert to the finer unit first. Raises as value does, its message naming Quantity.map2.
Algebra
add a b is map2 Nx.add a b. Its messages name Quantity.add, as in "Quantity.add: pix{b} does not convert to pix{a}: their quotient keeps pix{a}^-1 pix{b}" for a in pix{a} and b in pix{b}.
sub a b is map2 Nx.sub a b. Its messages name Quantity.sub.
mul a b is the product of the payloads in Unit.(unit a * unit b).
div a b is the quotient of the payloads, as Nx.div computes it, in Unit.(unit a / unit b).
pow n q is q's payload to the power n in Unit.(unit q ** n), computed by repeated products: integers wrap, and the exponent is exact whatever the payload's dtype, so pow 17 of -1. is -1. in float8. A negative n is the reciprocal of the power -n, and pow 0 is one, NaN included. Each product rounds, so a float result is within about n - 1 half-ulps of the exact power: pow 1000 in float32 can be hundreds of ulps off. Raises Invalid_argument "Quantity.pow: -1 is negative and the payload is int32" for a negative n on an integer payload.
root n q is the real n-th root of q's payload in Unit.root n (unit q), by Nx.pow at the exponent 1/n, then, for n >= 3, one Newton step that corrects the rounding of 1/n. A payload narrower than float32 is rooted in float32 and cast back, so the exponent is not rounded to its dtype. For odd n the root of an element below zero is the negative real root; for even n such an element raises Invalid_argument through Nx.check: at once on the host, at the next wait on another set, on a concrete payload or a traced one, as in "Quantity.root: element [3] of a float32 payload is below 0, whose root of order 2 is not real". NaN and -0. give what Nx.pow gives them. Raises Invalid_argument "Quantity.root: 0 is below 1" if n < 1.