"""
********************************************************************************
* Copyright (c) 2025 the Qrisp authors
*
* This program and the accompanying materials are made available under the
* terms of the Eclipse Public License 2.0 which is available at
* http://www.eclipse.org/legal/epl-2.0.
*
* This Source Code may also be made available under the following Secondary
* Licenses when the conditions for such availability set forth in the Eclipse
* Public License, v. 2.0 are satisfied: GNU General Public License, version 2
* with the GNU Classpath Exception which is
* available at https://www.gnu.org/software/classpath/license.html.
*
* SPDX-License-Identifier: EPL-2.0 OR GPL-2.0 WITH Classpath-exception-2.0
********************************************************************************
"""
import jax.numpy as jnp
import jax.lax as lax
import jax
from dataclasses import dataclass
from jax.tree_util import register_pytree_node_class
BASE = 2**32
DTYPE = jnp.uint32
BASE_FL = float(BASE)
[docs]
@register_pytree_node_class
@dataclass(frozen=True)
class BigInteger:
"""
Fixed-width, little-endian base-2^32 big integer for JAX.
This type represents non-negative integers using a fixed number of 32-bit
limbs (dtype=uint32) in little-endian order (digits[0] is the least significant
limb). It is compatible with JAX transformations (jit, vmap, pytrees).
All arithmetic (addition, subtraction, multiplication, shifts, division,
modulo, bitwise) is performed modulo 2^(32*n), where n = len(digits).
Overflow and underflow beyond the most-significant limb are discarded.
Notes
-----
- Width (number of limbs) is determined by `digits.shape[0]` and remains
constant; operations do not change the number of limbs.
- Operators assume both operands have the same number of limbs.
Attributes
----------
digits : jnp.ndarray
Little-endian limbs (dtype=uint32) of length `n`.
"""
digits: jnp.ndarray # Little-endian base-2^32
@jax.jit
def __call__(self):
"""
Return a float64 approximation of the integer value.
Computes sum_i digits[i] * (2^32)^i as float64. Exact only for values
that fit into float64; larger integers may lose precision.
Returns
-------
jnp.float64
Approximate numeric value (float64).
"""
r = lax.fori_loop(
0,
self.digits.shape[0],
lambda i, val: jnp.float64(
self.digits[i]) * BASE_FL ** jnp.float64(i) + val,
0.0,
)
return r
def tree_flatten(self):
"""
PyTree flatten for JAX.
Returns
-------
tuple
A pair `(children, aux_data)` where `children` is a tuple containing
the digits array, and `aux_data` is `None`.
"""
return (self.digits,), None
@classmethod
def tree_unflatten(cls, aux_data, children):
"""
PyTree unflatten for JAX.
Parameters
----------
aux_data : Any
Auxiliary data (unused, expected `None`).
children : tuple
Tuple containing the digits array.
Returns
-------
BigInteger
Reconstructed instance.
"""
return cls(*children)
[docs]
@staticmethod
def create_static(n, size):
"""
Create a BigInteger from Python using pure Python loops.
This variant does not use JAX primitives and is suitable for static
construction (e.g., outside `jit`). The result has exactly `size` limbs
with wraparound modulo 2^(32*size).
Parameters
----------
n : int or float
Non-negative number. Floats are truncated; very large floats
(> 2**53) may lose precision before conversion.
size : int
Number of limbs (digits) to allocate.
Returns
-------
BigInteger
Fixed-width representation of `n` modulo 2^(32*size).
"""
digits = []
for i in range(size):
digits.append(n % BASE)
n //= BASE
return BigInteger(jnp.array(digits, dtype=DTYPE))
[docs]
@staticmethod
def create(n, size):
"""
Create a BigInteger using JAX primitives.
Constructs a fixed-width BigInteger with exactly `size` limbs, interpreting
the input modulo 2^(32*size). JIT-friendly.
Parameters
----------
n : int or float or jnp.integer or jnp.floating
Non-negative number. Floats are truncated; very large floats
(> 2**53) may lose precision before conversion.
size : int
Number of limbs (digits) to allocate.
Returns
-------
BigInteger
Fixed-width representation of `n` modulo 2^(32*size).
Notes
-----
When called with a Python literal outside `jit`, the value must fit into
JAX's host integer range (typically up to 64 bits). For arbitrarily large
Python integers, prefer `create_static`.
"""
def body_fun(i, args):
digits, num = args
digits = digits.at[i].set(jnp.uint32(num % BASE))
num //= BASE
return digits, num
digits, _ = lax.fori_loop(0, size, body_fun, (jnp.zeros(size, dtype=DTYPE), n))
return BigInteger(digits)
@staticmethod
def create_dynamic(n, size):
"""
Alias of `create`.
Parameters
----------
n : int or float or jnp.integer or jnp.floating
Non-negative number.
size : int
Number of limbs (digits) to allocate.
Returns
-------
BigInteger
Fixed-width representation modulo 2^(32*size).
See Also
--------
BigInteger.create : JAX-compatible constructor.
"""
return BigInteger.create(n, size)
@jax.jit
def __add__(self, other: "BigInteger") -> "BigInteger":
"""
Add two BigIntegers with wraparound modulo 2^(32*n).
If `other` is a scalar, it is converted to the same width. The number
of limbs `n` is taken from `self`.
Parameters
----------
other : BigInteger or int
Addend.
Returns
-------
BigInteger
(self + other) mod 2^(32*n).
"""
n = self.digits.shape[0]
if not isinstance(other, BigInteger):
other = BigInteger.create(other, n)
a, b = self.digits, other.digits
result = jnp.zeros_like(a)
carry = jnp.uint64(0)
def add_step(i, state):
carry, result = state
s = jnp.uint64(a[i]) + jnp.uint64(b[i]) + carry
digit = jnp.uint32(s & (BASE - 1))
new_carry = jnp.uint64(s >> 32)
result = result.at[i].set(jnp.uint32(digit))
return new_carry, result
carry, result = lax.fori_loop(0, a.shape[0], add_step, (carry, result))
return BigInteger(result)
@jax.jit
def __sub__(self, other: "BigInteger") -> "BigInteger":
"""
Subtract two BigIntegers with wraparound modulo 2^(32*n).
If `other` is a scalar, it is converted to the same width. Computes
(self - other) mod 2^(32*n) using borrow propagation.
Parameters
----------
other : BigInteger or int
Subtrahend.
Returns
-------
BigInteger
(self - other) mod 2^(32*n).
"""
n = self.digits.shape[0]
if not isinstance(other, BigInteger):
other = BigInteger.create(other, n)
a, b = self.digits, other.digits
result = jnp.zeros_like(a)
def add_step(i, state):
carry, result = state
s, new_carry = lax.cond(
jnp.uint64(a[i]) >= jnp.uint64(b[i]) + carry,
lambda: (jnp.uint32(a[i] - b[i] - carry), 0),
lambda: (jnp.uint32(jnp.uint64(
a[i]) + jnp.uint64(BASE) - jnp.uint64(b[i]) - carry), 1),
)
result = result.at[i].set(jnp.uint32(s))
return new_carry, result
carry, result = lax.fori_loop(0, a.shape[0], add_step, (0, result))
return BigInteger(result)
@jax.jit
def __sub_alt__(self, other: "BigInteger") -> "BigInteger":
"""
Alternative subtraction using bitwise complement identity.
Computes (self - other) as `~((~self) + other)` with wraparound
modulo 2^(32*n).
Parameters
----------
other : BigInteger or int
Subtrahend.
Returns
-------
BigInteger
(self - other) mod 2^(32*n).
"""
return ~((~self) + other)
@jax.jit
def __mul__(self, other: "BigInteger") -> "BigInteger":
"""
Multiply two BigIntegers with wraparound modulo 2^(32*n).
Implements schoolbook multiplication and accumulates into `n` limbs,
discarding overflow beyond the n-th limb (wraparound).
Parameters
----------
other : BigInteger or int
Multiplier.
Returns
-------
BigInteger
(self * other) mod 2^(32*n).
"""
n = self.digits.shape[0]
if not isinstance(other, BigInteger):
other = BigInteger.create(other, n)
a, b = self.digits, other.digits
# result is uint32, but always promote to uint64 for arithmetic
result = jnp.zeros(n, dtype=DTYPE)
def outer_loop(i, result):
carry0 = jnp.uint64(0)
def inner_body(j, state):
res, carry = state
k = i + j
tmp = (
jnp.uint64(res[k])
+ jnp.uint64(a[i]) * jnp.uint64(b[j])
+ carry
)
new_digit = jnp.uint32(tmp & jnp.uint64(0xFFFFFFFF))
res = res.at[k].set(new_digit)
carry = tmp >> jnp.uint64(32)
return (res, carry)
# j runs where k = i + j < n -> j in [0, n - i)
result, _ = lax.fori_loop(0, n - i, inner_body, (result, carry0))
# Drop any remaining carry (mod base^n)
return result
result = lax.fori_loop(0, n, outer_loop, result)
return BigInteger(result)
@jax.jit
def __pow__(self, other):
"""
Integer exponentiation (square-and-multiply).
Performs `self ** other` by binary exponentiation, modulo 2^(32*n).
Parameters
----------
other : int or jnp.integer
Exponent (>= 0).
Returns
-------
BigInteger
self raised to power `other` modulo 2^(32*n).
Notes
-----
For `other == 0`, returns 1 (the multiplicative identity) with the same width.
"""
n = self.digits.shape[0]
base = self
exp = jnp.asarray(other, dtype=jnp.uint64)
acc = BigInteger.create(1, n)
def cond_fun(state):
base, exp, acc = state
return exp > 0
def body_fun(state):
base, exp, acc = state
acc = lax.cond((exp & jnp.uint64(1)) == jnp.uint64(1),
lambda a: a * base,
lambda a: a,
acc)
base = base * base
exp = exp >> jnp.uint64(1)
return base, exp, acc
_, _, acc = lax.while_loop(cond_fun, body_fun, (base, exp, acc))
return acc
def __repr__(self):
"""
String representation with limbs in little-endian order.
Returns
-------
str
String representation of digits (uint32 list).
"""
return self.digits.__repr__()
#return f"BigInteger(digits={self.digits.tolist()})"
@jax.jit
def __lt__(self, other: "BigInteger"):
"""
Less-than comparison between two fixed-width BigIntegers.
If `other` is a scalar, it is converted to the same width. Requires
both operands to have the same number of limbs.
Parameters
----------
other : BigInteger or int
Right-hand operand.
Returns
-------
jnp.bool_
True if `self < other` (unsigned), else False.
"""
n = self.digits.shape[0]
if not isinstance(other, BigInteger):
other = BigInteger.create(other, n)
m = other.digits.shape[0]
assert n == m
d0 = self.digits
d1 = other.digits
def body_fun(val):
i, res_found = val
i -= 1
res_found = lax.cond(d0[i] < d1[i], lambda: 1, lambda: -5)
res_found = lax.cond(d0[i] > d1[i], lambda: 0, lambda: res_found)
return i, res_found
def cond_fun(val):
i, res_found = val
return jnp.logical_and(res_found == -5, i > 0)
_, res = lax.while_loop(cond_fun, body_fun, (n, -5))
res = lax.cond(res == -5, lambda: 0, lambda: res)
return (res != 0)
@jax.jit
def __eq__(self, other: "BigInteger"):
"""
Equality comparison between two fixed-width BigIntegers.
If `other` is a scalar, it is converted to the same width.
Parameters
----------
other : BigInteger or int
Right-hand operand.
Returns
-------
jnp.bool_
True if all limbs are equal, else False.
"""
n = self.digits.shape[0]
if not isinstance(other, BigInteger):
other = BigInteger.create(other, n)
m = other.digits.shape[0]
assert n == m
d0 = self.digits
d1 = other.digits
return jnp.all(d0 == d1)
@jax.jit
def __ne__(self, other: "BigInteger"):
"""
Inequality comparison between two fixed-width BigIntegers.
Parameters
----------
other : BigInteger or int
Right-hand operand.
Returns
-------
jnp.bool_
True if any limb differs, else False.
"""
return jnp.logical_not(self == other)
@jax.jit
def __le__(self, other: "BigInteger"):
"""
Less-or-equal comparison between two fixed-width BigIntegers.
If `other` is a scalar, it is converted to the same width. Requires
both operands to have the same number of limbs.
Parameters
----------
other : BigInteger or int
Right-hand operand.
Returns
-------
jnp.bool_
True if `self <= other` (unsigned), else False.
"""
n = self.digits.shape[0]
if not isinstance(other, BigInteger):
other = BigInteger.create(other, n)
m = other.digits.shape[0]
assert n == m
d0 = self.digits
d1 = other.digits
def body_fun(val):
i, res_found = val
i -= 1
res_found = lax.cond(d0[i] < d1[i], lambda: 1, lambda: -5)
res_found = lax.cond(d0[i] > d1[i], lambda: 0, lambda: res_found)
return i, res_found
def cond_fun(val):
i, res_found = val
return jnp.logical_and(res_found == -5, i > 0)
_, res = lax.while_loop(cond_fun, body_fun, (n, -5))
res = lax.cond(res == -5, lambda: 1, lambda: res)
return (res != 0)
@jax.jit
def __lshift__(self, shift):
"""
Logical left shift by a non-negative number of bits.
Shifts bits left by `shift` and fills with zeros, within fixed width.
Bits shifted out of the most-significant end are discarded.
Parameters
----------
shift : int or jnp.integer
Number of bits to shift (>= 0).
Returns
-------
BigInteger
(self << shift) mod 2^(32*n).
"""
total_bits = self.digits.shape[0] * 32
zeros = BigInteger(jnp.zeros_like(self.digits))
def do_shift(_):
def body_fun(i, x):
return lax.cond(
self.get_bit(i) != 0,
lambda: x.flip_bit(i + shift),
lambda: x,
)
return lax.fori_loop(0, total_bits - shift, body_fun, zeros)
return lax.cond(
jnp.asarray(shift) >= total_bits,
lambda _: zeros,
do_shift,
operand=None,
)
@jax.jit
def __rshift__(self, shift):
"""
Logical right shift by a non-negative number of bits.
Shifts bits right by `shift` and fills with zeros, within fixed width.
Bits shifted out of the least-significant end are discarded.
Parameters
----------
shift : int or jnp.integer
Number of bits to shift (>= 0).
Returns
-------
BigInteger
(self >> shift) within fixed width.
"""
total_bits = self.digits.shape[0] * 32
zeros = BigInteger(jnp.zeros_like(self.digits))
def do_shift(_):
def body_fun(i, x):
return lax.cond(
self.get_bit(i) != 0,
lambda: x.flip_bit(i - shift),
lambda: x,
)
return lax.fori_loop(shift, total_bits, body_fun, zeros)
return lax.cond(
jnp.asarray(shift) >= total_bits,
lambda _: zeros,
do_shift,
operand=None,
)
@jax.jit
def __and__(self, other: "BigInteger") -> "BigInteger":
"""
Bitwise AND between two fixed-width BigIntegers.
If `other` is a scalar, it is converted to the same width.
Parameters
----------
other : BigInteger or int
Right-hand operand.
Returns
-------
BigInteger
self & other (limb-wise).
"""
n = self.digits.shape[0]
if not isinstance(other, BigInteger):
other = BigInteger.create(other, n)
return BigInteger(self.digits & other.digits)
@jax.jit
def __or__(self, other: "BigInteger") -> "BigInteger":
"""
Bitwise OR between two fixed-width BigIntegers.
If `other` is a scalar, it is converted to the same width.
Parameters
----------
other : BigInteger or int
Right-hand operand.
Returns
-------
BigInteger
self | other (limb-wise).
"""
n = self.digits.shape[0]
if not isinstance(other, BigInteger):
other = BigInteger.create(other, n)
return BigInteger(self.digits | other.digits)
@jax.jit
def __xor__(self, other: "BigInteger") -> "BigInteger":
"""
Bitwise XOR between two fixed-width BigIntegers.
If `other` is a scalar, it is converted to the same width.
Parameters
----------
other : BigInteger or int
Right-hand operand.
Returns
-------
BigInteger
self ^ other (limb-wise).
"""
n = self.digits.shape[0]
if not isinstance(other, BigInteger):
other = BigInteger.create(other, n)
return BigInteger(self.digits ^ other.digits)
@jax.jit
def __invert__(self) -> "BigInteger":
"""
Bitwise NOT on a fixed-width BigInteger.
Returns
-------
BigInteger
Bitwise complement of `self`, limb-wise.
"""
return BigInteger(~self.digits)
@jax.jit
def __mod__(self, other: "BigInteger"):
"""
Modulo operation `self % other` (fixed-width).
If `other` is a scalar, it is converted to the same width. Uses
`remainder_division`.
Parameters
----------
other : BigInteger or int
Modulus (must be non-zero).
Returns
-------
BigInteger
Remainder `r` with 0 <= r < other (when other != 0).
"""
if not isinstance(other, BigInteger):
other = BigInteger.create(other, self.digits.shape[0])
r, q = self.remainder_division(other)
return r
@jax.jit
def __floordiv__(self, other: "BigInteger"):
"""
Floor division `self // other` (fixed-width).
If `other` is a scalar, it is converted to the same width. Uses
`remainder_division`.
Parameters
----------
other : BigInteger or int
Divisor (must be non-zero).
Returns
-------
BigInteger
Quotient `q` such that self = other*q + r, 0 <= r < other.
"""
if not isinstance(other, BigInteger):
other = BigInteger.create(other, self.digits.shape[0])
r, q = self.remainder_division(other)
return q
@jax.jit
def get_bit(self, i: int):
"""
Get the value of the i-th bit (0-based, LSB=bit 0).
Parameters
----------
i : int or jnp.integer
Bit index (0 <= i < 32*n).
Returns
-------
jnp.uint32
Either 0 or 1.
"""
pos = i // 32
pos_in = i % 32
return (self.digits[pos] >> pos_in) & 1
@jax.jit
def flip_bit(self, i: int):
"""
Toggle the i-th bit (0-based, LSB=bit 0).
Parameters
----------
i : int or jnp.integer
Bit index to flip (0 <= i < 32*n).
Returns
-------
BigInteger
Copy of `self` with the specified bit toggled.
"""
pos = i // 32
pos_in = i % 32
ds = jnp.copy(self.digits)
ds = ds.at[pos].set(jnp.uint32(ds[pos] ^ (1 << pos_in)))
return BigInteger(ds)
@jax.jit
def bit_size(self):
"""
Return the position of the most significant set bit plus one.
Approximates the bit-length of the value: floor(log2(x)) + 1.
Returns 0 for zero.
Returns
-------
jnp.int64
Bit-length of the value (0 for zero).
"""
n = self.digits.shape[0]
is_zero = jnp.all(self.digits == 0)
def nonzero_len():
# Find index ms of the most-significant non-zero limb (scan from top)
def cond_fun(i):
return jnp.logical_and(i > 0, self.digits[i] == jnp.uint32(0))
def body_fun(i):
return i - 1
ms = lax.while_loop(cond_fun, body_fun, n - 1)
limb = self.digits[ms]
# bit length of a 32-bit limb
limb_bits = (jnp.floor(jnp.log2(jnp.float64(limb))).astype(jnp.int64) + 1)
return jnp.int64(32) * jnp.int64(ms) + limb_bits
return lax.cond(is_zero, lambda: jnp.int64(0), nonzero_len)
def remainder_division(self, other: "BigInteger"):
"""
Exact division using Knuth long division (base 2^32).
Computes quotient and remainder such that:
`self = other * q + r`, with `0 <= r < other`.
Parameters
----------
other : BigInteger
Divisor (must be non-zero), same width as `self`.
Returns
-------
tuple of BigInteger
`(r, q)` where `r` is the remainder and `q` is the quotient.
Notes
-----
Uses a normalized Knuth division with limb base 2^32. Both `self`
and `other` must have the same number of limbs.
"""
n = self.digits.shape[0]
m = other.digits.shape[0]
assert n == m
r_digits, q_digits = _remainder_division_knuth(self.digits, other.digits)
return BigInteger(r_digits), BigInteger(q_digits)
def get_larger(self):
"""
Given a BigInteger with n limbs, return a new BigInteger with 2n limbs and the
same number
Returns
-------
BigInteger
Larger BigInteger with 2n limbs
"""
pad = jnp.zeros(self.digits.shape[0], dtype=self.digits.dtype)
return BigInteger(jnp.concatenate([self.digits, pad], axis=0))
@jax.jit
def _clz32(x):
"""
Count leading zeros in a 32-bit word.
Parameters
----------
x : jnp.uint32
Input 32-bit word.
Returns
-------
jnp.int32
Number of leading zeros in `x` (0..32).
"""
# Count leading zeros in a 32-bit word (simple loop)
def cond_fun(state):
x, s = state
return jnp.logical_and(s < jnp.int32(32), (x & jnp.uint32(0x80000000)) == jnp.uint32(0))
def body_fun(state):
x, s = state
return jnp.uint32(x << jnp.uint32(1)), s + jnp.int32(1)
_, s = lax.while_loop(cond_fun, body_fun, (x, jnp.int32(0)))
return s
@jax.jit
def _ms_length(a: jnp.ndarray):
"""
Effective limb length (highest non-zero index + 1).
Parameters
----------
a : jnp.ndarray
Little-endian uint32 limb array.
Returns
-------
jnp.int32
Effective length in limbs (0 if all zero).
"""
# Effective limb length: highest non-zero index + 1, or 0 if all zero.
n = a.shape[0]
all_zero = jnp.all(a == jnp.uint32(0))
def find_ms():
def cond_fun(i):
return jnp.logical_and(i > 0, a[i] == jnp.uint32(0))
def body_fun(i):
return i - 1
ms = lax.while_loop(cond_fun, body_fun, jnp.int32(n - 1))
return ms + jnp.int32(1)
return lax.cond(all_zero, lambda: jnp.int32(0), find_ms)
@jax.jit
def _shl_bits(arr: jnp.ndarray, s):
"""
Shift-left by `s` bits across limbs (0 <= s < 32).
Parameters
----------
arr : jnp.ndarray
Little-endian uint32 limb array to shift.
s : jnp.int32
Bit count (0..31).
Returns
-------
tuple
`(out, carry_out)` where `out` is the shifted array and `carry_out`
is the carry from the most-significant limb (uint32).
"""
# Shift-left by s bits (0<=s<32) across limbs (little endian).
s_u = jnp.uint32(s)
n = arr.shape[0]
def no_shift():
return jnp.array(arr, dtype=jnp.uint32), jnp.uint32(0)
def do_shift():
out = jnp.zeros_like(arr)
carry = jnp.uint32(0)
def body(i, state):
carry, out = state
ai = arr[i]
low = jnp.uint32((jnp.uint64(ai) << jnp.uint64(s_u))
& jnp.uint64(0xFFFFFFFF))
new_digit = jnp.uint32(low | carry)
out = out.at[i].set(new_digit)
# IMPORTANT: parentheses for precedence
new_carry = jnp.uint32(jnp.uint64(ai) >> (jnp.uint64(32) - jnp.uint64(s_u)))
return new_carry, out
carry_out, out = lax.fori_loop(0, n, body, (carry, out))
return out, carry_out
return lax.cond(s == jnp.int32(0), no_shift, do_shift)
@jax.jit
def _shr_bits(arr: jnp.ndarray, s) -> jnp.ndarray:
"""
Shift-right by `s` bits across limbs (0 <= s < 32).
Parameters
----------
arr : jnp.ndarray
Little-endian uint32 limb array to shift.
s : jnp.int32
Bit count (0..31).
Returns
-------
jnp.ndarray
Shifted array (little-endian, uint32).
"""
# Shift-right by s bits (0<=s<32) across limbs (little endian).
s_u = jnp.uint32(s)
n = arr.shape[0]
def no_shift():
return jnp.array(arr, dtype=jnp.uint32)
def do_shift():
out = jnp.zeros_like(arr)
carry = jnp.uint32(0)
def body(i, state):
carry, out = state
idx = n - 1 - i
ai = arr[idx]
high = jnp.uint32((jnp.uint64(carry) << (jnp.uint64(
32) - jnp.uint64(s_u))) & jnp.uint64(0xFFFFFFFF))
new_digit = jnp.uint32(
(jnp.uint64(ai) >> jnp.uint64(s_u)) | jnp.uint64(high))
out = out.at[idx].set(new_digit)
new_carry = jnp.uint32(ai & (jnp.uint32(1) << s_u) - jnp.uint32(1))
return new_carry, out
_, out = lax.fori_loop(0, n, body, (carry, out))
return out
return lax.cond(s == jnp.int32(0), no_shift, do_shift)
@jax.jit
def _divmod_single_limb(u: jnp.ndarray, d):
"""
Divide a multi-limb number by a single 32-bit limb.
Parameters
----------
u : jnp.ndarray
Little-endian uint32 limb array (dividend).
d : jnp.uint32
Single-limb divisor (must be non-zero).
Returns
-------
tuple
`(q_digits, r_low)` where `q_digits` is the quotient array (uint32)
and `r_low` is the remainder (uint32).
"""
# Divide multi-limb u by single 32-bit d (d != 0). Returns (q_digits, r_low).
n = u.shape[0]
q = jnp.zeros_like(u)
rem = jnp.uint64(0)
def body(i, state):
q, rem = state
idx = n - 1 - i
cur = (rem << jnp.uint64(32)) + jnp.uint64(u[idx])
qi = jnp.uint32(cur // jnp.uint64(d))
rem = jnp.uint64(cur % jnp.uint64(d))
q = q.at[idx].set(qi)
return q, rem
q, rem = lax.fori_loop(0, n, body, (q, rem))
return q, jnp.uint32(rem)
@jax.jit
def _remainder_division_knuth(u: jnp.ndarray, v: jnp.ndarray) -> tuple[jnp.ndarray, jnp.ndarray]:
"""
Knuth long division (base 2^32) for arrays of equal length.
Parameters
----------
u : jnp.ndarray
Dividend limbs (little-endian uint32).
v : jnp.ndarray
Divisor limbs (little-endian uint32), not all zero.
Returns
-------
tuple of jnp.ndarray
`(r_digits, q_digits)` where both are little-endian uint32 arrays
of the same length as `u`. The relation `u = v * q + r` holds with
`0 <= r < v`.
Notes
-----
Normalizes the divisor and uses a classic multi-precision division scheme.
"""
# Knuth long division base 2^32. Inputs uint32 arrays (same length N).
# Returns (r_digits, q_digits), both length N.
N = u.shape[0]
m = _ms_length(v) # effective length of divisor
n_eff = _ms_length(u) # effective length of dividend
# Divisor zero => return (0, 0)
def div_by_zero():
return jnp.zeros_like(u), jnp.zeros_like(u)
def normal_path():
# If u < v: q = 0, r = u
def less_path():
return u, jnp.zeros_like(u)
def ge_path():
# Single-limb fast path
def single_limb_path():
d = v[0]
q, r_low = _divmod_single_limb(u, d)
r = jnp.zeros_like(u).at[0].set(r_low)
return r, q
# Multi-limb path (m >= 2)
def multi_limb_path():
ms_idx = m - jnp.int32(1)
v_msw = v[ms_idx]
s = _clz32(v_msw) # 0..31
v_norm, _ = _shl_bits(v, s) # length N
u_norm_part, carry_u = _shl_bits(u, s) # length N, carry
u_norm = jnp.concatenate([u_norm_part, jnp.array(
[carry_u], dtype=jnp.uint32)], axis=0) # N+1
q = jnp.zeros_like(u)
# how many quotient positions we fill
j_count = jnp.maximum(n_eff - m, jnp.int32(0))
BASE_MASK = jnp.uint64(0xFFFFFFFF)
BASE64 = jnp.uint64(1) << jnp.uint64(32)
def body(t, state):
u_norm, q = state
active = t <= j_count
def do_step(state):
u_norm, q = state
j = j_count - t # 0..j_count
ujm = j + m # index to top limb of current window
# Estimate qhat from top two limbs of u_norm and top limb of v_norm
u2 = jnp.uint64(u_norm[ujm])
u1 = jnp.uint64(u_norm[ujm - 1])
num = (u2 << jnp.uint64(32)) + u1
den = jnp.uint64(v_norm[m - 1])
qhat = jnp.minimum(num // den, jnp.uint64(0xFFFFFFFF))
rhat = num - qhat * den
# Refinement (at most twice)
vm2 = jnp.uint64(v_norm[m - 2])
u0 = jnp.uint64(u_norm[ujm - 2])
cond1 = jnp.logical_and(qhat == jnp.uint64(0xFFFFFFFF),
qhat * vm2 > (rhat << jnp.uint64(32)) + u0)
qhat = jnp.where(cond1, qhat - jnp.uint64(1), qhat)
rhat = jnp.where(cond1, rhat + den, rhat)
cond2 = qhat * vm2 > (rhat << jnp.uint64(32)) + u0
qhat = jnp.where(cond2, qhat - jnp.uint64(1), qhat)
rhat = jnp.where(cond2, rhat + den, rhat)
# Subtract qhat * v_norm from u_norm[j ... j+m]
carry = jnp.uint64(0)
def sub_body(i, sstate):
u_norm, carry = sstate
vi = jnp.uint64(v_norm[i])
p = qhat * vi + carry
p_low = jnp.uint32(p & BASE_MASK)
p_high = p >> jnp.uint64(32)
idx = j + i
u_di = u_norm[idx]
borrow = u_di < p_low
new_u = jnp.uint32(
(jnp.uint64(u_di) + BASE64 - jnp.uint64(p_low)) & BASE_MASK)
u_norm = u_norm.at[idx].set(new_u)
carry = p_high + jnp.uint64(borrow)
return u_norm, carry
u_norm, carry = lax.fori_loop(0, m, sub_body, (u_norm, carry))
top_before = u_norm[ujm]
underflow = top_before < jnp.uint32(carry)
top_after = jnp.uint32(
(jnp.uint64(top_before) + BASE64 - jnp.uint64(carry)) & BASE_MASK)
u_norm = u_norm.at[ujm].set(top_after)
def fix_underflow(uq_state):
u_norm, q, qhat, j = uq_state
qhat = qhat - jnp.uint64(1)
carry2 = jnp.uint64(0)
def add_body(i, astate):
u_norm, carry2 = astate
idx = j + i
ssum = jnp.uint64(u_norm[idx]) + \
jnp.uint64(v_norm[i]) + carry2
new_digit = jnp.uint32(ssum & BASE_MASK)
carry2 = ssum >> jnp.uint64(32)
u_norm = u_norm.at[idx].set(new_digit)
return u_norm, carry2
u_norm, carry2 = lax.fori_loop(
0, m, add_body, (u_norm, carry2))
u_norm = u_norm.at[ujm].set(jnp.uint32(
jnp.uint64(u_norm[ujm]) + jnp.uint64(1)))
return u_norm, q, qhat, j
u_norm, q, qhat, _ = lax.cond(
underflow,
fix_underflow,
lambda uq_state: uq_state,
operand=(u_norm, q, qhat, j),
)
q = q.at[j].set(jnp.uint32(qhat & BASE_MASK))
return u_norm, q
u_norm, q = lax.cond(
active, do_step, lambda state: state, operand=(u_norm, q))
return u_norm, q
u_norm, q = lax.fori_loop(0, N, body, (u_norm, q))
# Build r_full by copying the first m limbs (dynamic) of u_norm, then de-normalize
r_norm_full = jnp.zeros_like(u)
def copy_body(i, r_acc):
val = jnp.where(i < m, u_norm[i], jnp.uint32(0))
return r_acc.at[i].set(val)
r_norm_full = lax.fori_loop(0, N, copy_body, r_norm_full)
r = _shr_bits(r_norm_full, s)
return r, q
return lax.cond(m == jnp.int32(1), single_limb_path, multi_limb_path)
return lax.cond(n_eff < m, less_path, ge_path)
return lax.cond(m == jnp.int32(0), div_by_zero, normal_path)
@jax.jit
def bi_modinv(a: BigInteger, m: BigInteger) -> BigInteger:
"""
Modular inverse using an Extended Euclidean Algorithm variant.
Finds `t` such that `(a * t) % m == 1`, assuming `gcd(a, m) == 1`.
Returns the non-negative representative in `[0, m)`.
Parameters
----------
a : BigInteger
Value to invert (same width as `m`).
m : BigInteger
Modulus (must be > 0; often odd in Montgomery contexts).
Returns
-------
BigInteger
`t = a^{-1} mod m`.
Raises
------
AssertionError
If inputs do not have the same width.
Notes
-----
Uses `//` and `%` which rely on exact multi-precision division implemented
in this module. Both `a` and `m` are treated as fixed-width unsigned values.
"""
# Widen to 2n limbs
n = a.digits.shape[0]
pad = jnp.zeros(n, dtype=DTYPE)
a = BigInteger(jnp.concatenate([a.digits, pad], axis=0))
m = BigInteger(jnp.concatenate([m.digits, pad], axis=0))
bi0 = BigInteger.create(0, 2*n)
bi1 = BigInteger.create(1, 2*n)
t, new_t = bi0, bi1
r, new_r = m, a
def cond(state):
_, new_t, r, new_r = state
return new_r != bi0
def body(state):
t, new_t, r, new_r = state
quotient = r // new_r
# Unsigned arithmetic with wraparound simulated modulo behavior
t_updated = (t + m - (quotient * new_t) % m) % m
r_updated = r - quotient * new_r
return new_t, t_updated, new_r, r_updated
final_t, _, final_r, _ = lax.while_loop(cond, body, (t, new_t, r, new_r))
return BigInteger(final_t.digits[0:n])
@jax.jit
def bi_extended_euclidean(a, b):
"""
Extended Euclidean Algorithm (fixed-width arithmetic).
Computes `g, x, y` such that `a*x + b*y = g = gcd(a, b)`. Compatible with
`jax.jit`. Inputs are expected to be `BigInteger` of the same width (or
scalars that are promoted to that width).
Parameters
----------
a : BigInteger or int
First number.
b : BigInteger or int
Second number.
Returns
-------
tuple
`(g, x, y)` where `g` is the gcd and `x, y` are Bézout coefficients
(all `BigInteger`).
Notes
-----
All operations occur modulo 2^(32*n), so the interpretation follows
fixed-width arithmetic semantics.
"""
n = a.digits.shape[0]
if not isinstance(b, BigInteger):
b = BigInteger.create(b, n)
bi0 = BigInteger.create(0, n)
bi1 = BigInteger.create(1, n)
# State: (r, old_r, s, old_s, t, old_t)
state = (b, a, bi0, bi1, bi1, bi0)
def cond_fun(state):
r, old_r, s, old_s, t, old_t = state
return r != bi0
def body_fun(state):
r, old_r, s, old_s, t, old_t = state
quotient = old_r // r
return (
old_r - quotient * r,
r,
old_s - quotient * s,
s,
old_t - quotient * t,
t,
)
r, old_r, s, old_s, t, old_t = lax.while_loop(cond_fun, body_fun, state)
# old_r is gcd, old_s and old_t are the Bézout coefficients
return old_r, old_s, old_t
@jax.jit
def bi_montgomery_encode(x: BigInteger, R: BigInteger, modulus: BigInteger) -> BigInteger:
"""
Montgomery encode: map x to (x * R) mod modulus, without intermediate wraparound.
This routine widens all operands to 2n limbs (where n is the current width) to compute
the product exactly, reduces the 2n-limb product modulo the widened modulus (also 2n
limbs), and then shrinks the remainder back to n limbs. This avoids the truncation
that would occur if the n-limb product were computed directly (since BigInteger
arithmetic is modulo 2^(32*n)).
Parameters
----------
x : BigInteger
Value to encode (width n limbs).
R : BigInteger
Montgomery radix (width n limbs), typically R = 2^m in the same limb system or
equivalent precomputed value for the chosen modulus.
modulus : BigInteger
Modulus (width n limbs). In radix-2 Montgomery arithmetic, modulus is assumed odd.
Returns
-------
BigInteger
x in Montgomery form, i.e., (x * R) mod modulus, represented on n limbs.
Notes
-----
- The 2n-limb widening ensures that (x * R) is computed without loss before the
modulo reduction. The remainder is strictly less than modulus < 2^(32*n), so it
fits in n limbs and can be safely truncated back.
- All three inputs must share the same limb width n.
"""
n = modulus.digits.shape[0]
# Widen to 2n limbs
pad = jnp.zeros(n, dtype=DTYPE)
x2 = BigInteger(jnp.concatenate([x.digits, pad], axis=0))
R2 = BigInteger(jnp.concatenate([R.digits, pad], axis=0))
N2 = BigInteger(jnp.concatenate([modulus.digits, pad], axis=0))
# Exact product at 2n limbs, reduce modulo widened modulus, then shrink to n limbs
prod2 = x2 * R2
rem2 = prod2 % N2
return BigInteger(rem2.digits[:n])
@jax.jit
def bi_montgomery_decode(x_mon: BigInteger, R: BigInteger, modulus: BigInteger) -> BigInteger:
"""
Montgomery decode: map x_mon to (x_mon * R^{-1}) mod modulus, without wraparound.
This routine computes invR = R^{-1} mod modulus at width n, widens x_mon and invR
to 2n limbs, performs the 2n-limb product and reduction modulo the widened modulus,
and then shrinks the remainder back to n limbs. This avoids truncation during the
multiply step in fixed-width arithmetic.
Parameters
----------
x_mon : BigInteger
Value in Montgomery form (width n limbs).
R : BigInteger
Montgomery radix used for encoding (width n limbs).
modulus : BigInteger
Modulus (width n limbs). In radix-2 Montgomery arithmetic, modulus is assumed odd.
Returns
-------
BigInteger
Decoded value, i.e., (x_mon * R^{-1}) mod modulus, represented on n limbs.
Notes
-----
- The 2n-limb widening ensures that (x_mon * invR) is computed exactly before the
modulo reduction. The remainder is < modulus < 2^(32*n), so it fits within n limbs.
- All three inputs must share the same limb width n.
"""
n = modulus.digits.shape[0]
# Compute R^{-1} mod modulus at n limbs
invR = bi_modinv(R, modulus)
# Widen operands and modulus to 2n limbs
pad = jnp.zeros(n, dtype=DTYPE)
x2 = BigInteger(jnp.concatenate([x_mon.digits, pad], axis=0))
inv2 = BigInteger(jnp.concatenate([invR.digits, pad], axis=0))
N2 = BigInteger(jnp.concatenate([modulus.digits, pad], axis=0))
# Exact product at 2n limbs, reduce modulo widened modulus, then shrink to n limbs
prod2 = x2 * inv2
rem2 = prod2 % N2
return BigInteger(rem2.digits[:n])
@jax.jit
def _bi_all_ones(n_limbs: int) -> BigInteger:
"""
Internal: BigInteger with all digits set to 0xFFFFFFFF (acts like +infinity bound).
Parameters
----------
n_limbs : int
Number of 32-bit limbs.
Returns
-------
BigInteger
All-ones value, width n_limbs.
"""
return BigInteger(jnp.full((n_limbs,), jnp.uint32(0xFFFFFFFF), dtype=DTYPE))
def bi_contfrac_best_approx(a: BigInteger,
b: BigInteger,
max_den: BigInteger | None = None,
max_iters: int | None = None) -> tuple[BigInteger, BigInteger]:
"""
Best rational approximation p/q to a/b via continued fractions, with q <= max_den.
This computes a convergent (or intermediate convergent) of the continued fraction
expansion of a/b using the standard recurrence, stopping when either:
- remainder becomes 0 (exact rational found), or
- the next denominator would exceed `max_den`; in that case the best
"intermediate" convergent within the bound is returned using
t = floor((max_den - q0)/q1), p = t*p1 + p0, q = t*q1 + q0.
The implementation uses a single JAX while_loop; it allocates no dynamic lists
and is suitable for very large BigIntegers and JIT/vmap usage.
Parameters
----------
a : BigInteger
Numerator (non-negative). Width must match `b`.
b : BigInteger
Denominator (non-zero, non-negative). Width must match `a`.
max_den : BigInteger or None, optional
Upper bound on the denominator q. If None, a built-in all-ones
value of the same width is used (effectively unbounded).
max_iters : int or None, optional
Maximum number of continued-fraction steps (safety cap). If None,
defaults to 2*bitlen + 4, where bitlen = 32 * n_limbs.
Returns
-------
(BigInteger, BigInteger)
A pair (p, q) with q > 0, p/q the best approximation to a/b
under the bound q <= max_den. For exact rationals, returns the exact p/q.
Raises
------
AssertionError
If widths of a and b differ, or if b == 0.
Notes
-----
- All operands must share the same width (number of limbs).
- Arithmetic is exact in the fixed-width ring; for valid inputs arising in Shor's
post-processing (with q <= b and typical bounds <= b), no wrap-around occurs.
"""
n_limbs = a.digits.shape[0]
assert n_limbs == b.digits.shape[0], "BigInteger widths must match"
bi0 = BigInteger.create(0, n_limbs)
bi1 = BigInteger.create(1, n_limbs)
assert b != bi0, "Denominator b must be non-zero"
# Prepare bound and iteration cap
if max_den is None:
max_den = _bi_all_ones(n_limbs)
else:
assert max_den.digits.shape[0] == n_limbs, "max_den must match width of a,b"
if max_iters is None:
max_iters = int(2 * 32 * n_limbs + 4)
@jax.jit
def _loop(a0: BigInteger, b0: BigInteger,
p0: BigInteger, p1: BigInteger,
q0: BigInteger, q1: BigInteger,
res_p: BigInteger, res_q: BigInteger,
done: bool, i: int) -> tuple:
"""
One step of CF with bound handling. Internal helper for while_loop.
"""
# Compute quotient and remainder
quot = a0 // b0
rem = a0 % b0
# Next convergent (pn/qn)
pn = quot * p1 + p0
qn = quot * q1 + q0
# Conditions
exact_end = (rem == bi0)
# exceed if next denominator would be > max_den
exceed = jnp.logical_not(qn <= max_den)
# Intermediate convergent if exceed: t = floor((max_den - q0)/q1)
# Guard q1==0 (only possible at very first step if max_den < 1)
q1_is_zero = (q1 == bi0)
t_num = (max_den - q0)
t = t_num // jax.lax.cond(q1_is_zero, lambda: bi1, lambda: q1)
p_bound = t * p1 + p0
q_bound = t * q1 + q0
# Choose final result when done at this step
sel_p = jax.lax.select(exceed, p_bound, pn)
sel_q = jax.lax.select(exceed, q_bound, qn)
# Update result slots if we finish now
res_p_next = jax.lax.select(jnp.logical_or(exceed, exact_end), sel_p, res_p)
res_q_next = jax.lax.select(jnp.logical_or(exceed, exact_end), sel_q, res_q)
done_next = jnp.logical_or(done, jnp.logical_or(exceed, exact_end))
# Prepare state for next iteration (if not done)
a_next = jax.lax.select(done_next, a0, b0)
b_next = jax.lax.select(done_next, b0, rem)
p0_next = jax.lax.select(done_next, p0, p1)
p1_next = jax.lax.select(done_next, p1, pn)
q0_next = jax.lax.select(done_next, q0, q1)
q1_next = jax.lax.select(done_next, q1, qn)
i_next = i + jnp.int32(1)
return (a_next, b_next, p0_next, p1_next, q0_next, q1_next,
res_p_next, res_q_next, done_next, i_next)
def cond_fn(state):
a0, b0, p0, p1, q0, q1, res_p, res_q, done, i = state
under_cap = (i < jnp.int32(max_iters))
return jnp.logical_and(jnp.logical_not(done), under_cap)
# Initial CF state: p[-2]=0, p[-1]=1; q[-2]=1, q[-1]=0
init_state = (a, b, bi0, bi1, bi1, bi0, bi0, bi1, jnp.bool_(False), jnp.int32(0))
final_state = lax.while_loop(cond_fn, _loop, init_state)
_, _, _, _, _, _, out_p, out_q, _, _ = final_state
return out_p, out_q
def bi_shor_recover_denominator(a: BigInteger,
b: BigInteger,
N_bound: BigInteger | int,
max_iters: int | None = None) -> BigInteger:
"""
Recover the candidate period denominator r from a/b for Shor's algorithm.
This computes the best convergent p/q of a/b with q <= N_bound, where N_bound
is typically the number to factor (or a small multiple), returning q.
Parameters
----------
a : BigInteger
Numerator of measurement ratio (0 <= a < b).
b : BigInteger
Denominator (typically a power of two, b = 2^t).
N_bound : BigInteger or int
Upper bound for the denominator r (e.g., N).
max_iters : int or None, optional
Maximum CF steps (safety cap). See `bi_contfrac_best_approx`.
Returns
-------
BigInteger
The recovered denominator r (candidate period). Use with standard
Shor post-checks (e.g., verify a close approximation and test that
x^r ≡ 1 mod N, handle r even, etc.).
Notes
-----
- If N_bound is an int, it is promoted to a BigInteger of the same width as a,b.
- Returns the denominator of the CF-derived convergent; caller should perform
the usual validity checks for Shor (closeness, non-trivial factor conditions, etc.).
"""
n_limbs = a.digits.shape[0]
if isinstance(N_bound, BigInteger):
max_den = N_bound
assert max_den.digits.shape[0] == n_limbs, "N_bound width must match a,b"
else:
max_den = BigInteger.create(int(N_bound), n_limbs)
p, q = bi_contfrac_best_approx(a, b, max_den=max_den, max_iters=max_iters)
return q
def bi_contfrac_convergents(a: BigInteger,
b: BigInteger,
max_terms: int | None = None):
"""
Generator of convergents p/q for a/b (Python generator; not JAX-traced).
Yields successive convergents (p_k, q_k) using the standard recurrence until
remainder becomes 0 or `max_terms` is reached. Useful for debugging or
non-jitted workflows. For high-performance in-jit use, prefer
`bi_contfrac_best_approx`.
Parameters
----------
a : BigInteger
Numerator.
b : BigInteger
Denominator (non-zero).
max_terms : int or None, optional
Maximum number of convergents to yield.
Yields
------
(BigInteger, BigInteger)
Convergent pairs (p_k, q_k), k = 0, 1, 2, ...
"""
n_limbs = a.digits.shape[0]
assert n_limbs == b.digits.shape[0], "BigInteger widths must match"
bi0 = BigInteger.create(0, n_limbs)
bi1 = BigInteger.create(1, n_limbs)
assert b != bi0, "Denominator b must be non-zero"
a0, b0 = a, b
p0, p1 = bi0, bi1
q0, q1 = bi1, bi0
k = 0
while True:
q = a0 // b0
r = a0 % b0
pn = q * p1 + p0
qn = q * q1 + q0
yield pn, qn
if r == bi0:
break
a0, b0 = b0, r
p0, p1 = p1, pn
q0, q1 = q1, qn
k += 1
if max_terms is not None and k >= max_terms:
break
