Coverage for src/bartz/_jaxext/_jaxext.py: 97%
151 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-25 11:03 +0000
1# bartz/src/bartz/_jaxext/_jaxext.py
2#
3# Copyright (c) 2024-2026, The Bartz Contributors
4#
5# This file is part of bartz.
6#
7# Permission is hereby granted, free of charge, to any person obtaining a copy
8# of this software and associated documentation files (the "Software"), to deal
9# in the Software without restriction, including without limitation the rights
10# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
11# copies of the Software, and to permit persons to whom the Software is
12# furnished to do so, subject to the following conditions:
13#
14# The above copyright notice and this permission notice shall be included in all
15# copies or substantial portions of the Software.
16#
17# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
18# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
19# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
20# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
21# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
22# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
23# SOFTWARE.
25"""Implementation of miscellaneous jax extension utilities."""
27import math
28import sys
29from collections.abc import Callable, Generator, Sequence
30from contextlib import contextmanager
31from functools import partial, wraps
32from typing import Any, ParamSpec, TypeVar
34import jax
35from jax import Device, ensure_compile_time_eval, lax, random, shard_map, tree, vmap
36from jax import numpy as jnp
37from jax.dtypes import prng_key
38from jax.scipy.special import ndtr, ndtri
39from jax.sharding import PartitionSpec
40from jax.typing import DTypeLike
41from jaxtyping import (
42 Array,
43 Bool,
44 Float32,
45 Integer,
46 Key,
47 PyTree,
48 Scalar,
49 ScalarLike,
50 Shaped,
51)
52from jaxtyping import config as jaxtyping_config
54from bartz._jaxext._jit import jit
56if sys.version_info >= (3, 13):
57 from typing import TypeIs
58else: # WORKAROUND(python<3.13): typing.TypeIs was added in 3.13
59 from typing_extensions import TypeIs
61_P = ParamSpec('_P')
62_R = TypeVar('_R')
65@contextmanager
66def jaxtyping_disabled() -> Generator[None, None, None]:
67 """Temporarily disable jaxtyping runtime type-checking.
69 This also disables `beartype`, because the jaxtyping import hook applies it
70 as ``jaxtyped(typechecker=beartype)`` and `jaxtyped` short-circuits to the
71 undecorated function when type-checking is disabled. Used to park
72 deliberately wrong-typed intermediates (e.g. `_LazyArray` leaves) in an
73 `equinox.Module` during construction.
74 """
75 old = jaxtyping_config.jaxtyping_disable
76 jaxtyping_config.update('jaxtyping_disable', True)
77 try:
78 yield
79 finally:
80 jaxtyping_config.update('jaxtyping_disable', old)
83def float32_matmuls(fun: Callable[_P, _R]) -> Callable[_P, _R]:
84 """Run/trace `fun` under full-float32 matmul precision."""
86 @wraps(fun)
87 def wrapper(*args: _P.args, **kw: _P.kwargs) -> _R:
88 with jax.default_matmul_precision('float32'):
89 return fun(*args, **kw)
91 return wrapper
94def vmap_nodoc(fun: Callable, *args: Any, **kw: Any) -> Callable:
95 """
96 Acts like `jax.vmap` but preserves the docstring of the function unchanged.
98 This is useful if the docstring already takes into account that the
99 arguments have additional axes due to vmap.
100 """
101 doc = fun.__doc__
102 fun = vmap(fun, *args, **kw)
103 fun.__doc__ = doc
104 return fun
107def sliced_map(
108 f: Callable[[PyTree[Array, ' I']], PyTree[Array, ' O']],
109 xs: PyTree[Array, ' I'],
110 *,
111 batch_size: int,
112) -> PyTree[Array, ' O']:
113 """
114 Like `jax.lax.map` with `batch_size`, but read `xs` by slicing.
116 The scan body slices each batch out of the closed-over `xs` with
117 `lax.dynamic_slice_in_dim` instead of consuming `xs` reshaped into scan xs.
118 Under `vmap`, the batching rule for values closed over by a scan moves the
119 batch axes to the front, so leading batch axes in the leaves of `xs` are
120 consumed in place, while `lax.map` would transpose them to sit after the
121 scanned and `batch_size` axes.
123 Parameters
124 ----------
125 f
126 The function to apply elementwise, taking and returning pytrees of
127 arrays.
128 xs
129 A pytree of arrays to map over along their first axes.
130 batch_size
131 The number of elements to process at a time, vectorized. The remainder
132 of the division of the number of elements by `batch_size` is processed
133 as one additional smaller batch.
135 Returns
136 -------
137 A pytree of stacked outputs of `f`, like `jax.lax.map`.
138 """
139 (num_el,) = {x.shape[0] for x in tree.leaves(xs)}
141 # shortcut if no batching needed
142 if batch_size >= num_el:
143 return vmap(f)(xs)
145 num_batches, remainder = divmod(num_el, batch_size)
147 def apply_batch(start: Integer[Array, ''] | int, size: int) -> PyTree[Array, ' O']:
148 def slice_batch(
149 x: Shaped[Array, ' num_el *shape'],
150 ) -> Shaped[Array, ' size *shape']:
151 return lax.dynamic_slice_in_dim(x, start, size, axis=0)
153 return vmap(f)(tree.map(slice_batch, xs))
155 def loop(_: None, start: Integer[Array, '']) -> tuple[None, PyTree[Array, ' O']]:
156 return None, apply_batch(start, batch_size)
158 _, out = lax.scan(loop, None, batch_size * jnp.arange(num_batches))
159 out = tree.map(lambda x: x.reshape(num_batches * batch_size, *x.shape[2:]), out)
160 if remainder:
161 rest = apply_batch(num_batches * batch_size, remainder)
162 out = tree.map(lambda x, y: jnp.concatenate([x, y], axis=0), out, rest)
163 return out
166def minimal_unsigned_dtype(value: int) -> DTypeLike:
167 """Return the smallest unsigned integer dtype that can represent `value`."""
168 if value < 2**8:
169 return jnp.uint8
170 if value < 2**16: 170 ↛ 172line 170 didn't jump to line 172 because the condition on line 170 was always true
171 return jnp.uint16
172 if value < 2**32:
173 return jnp.uint32
174 return jnp.uint64
177@jit(static_argnums=(1,))
178def unique(
179 x: Shaped[Array, ' _'], size: int, fill_value: ScalarLike
180) -> tuple[Shaped[Array, ' {size}'], int | Integer[Array, '']]:
181 """
182 Restricted version of `jax.numpy.unique` that uses less memory.
184 Parameters
185 ----------
186 x
187 The input array.
188 size
189 The length of the output.
190 fill_value
191 The value to fill the output with if `size` is greater than the number
192 of unique values in `x`.
194 Returns
195 -------
196 out : Shaped[Array, '{size}']
197 The unique values in `x`, sorted, and right-padded with `fill_value`.
198 actual_length : int
199 The number of used values in `out`.
200 """
201 if x.size == 0:
202 return jnp.full(size, fill_value, x.dtype), 0
203 if size == 0:
204 return jnp.empty(0, x.dtype), 0
205 x = jnp.sort(x)
207 def loop(
208 carry: tuple[Scalar, Scalar, Shaped[Array, ' size']], x: Scalar
209 ) -> tuple[tuple[Scalar, Scalar, Shaped[Array, ' size']], None]:
210 i_out, last, out = carry
211 i_out = jnp.where(x == last, i_out, i_out + 1)
212 out = out.at[i_out].set(x)
213 return (i_out, x, out), None
215 carry = jnp.array(0), x[0], jnp.full(size, fill_value, x.dtype)
217 def run(unroll: int) -> tuple[Shaped[Array, ' size'], Scalar]:
218 (actual_length, _, out), _ = lax.scan(loop, carry, x[:size], unroll=unroll)
219 return out, actual_length + 1
221 # The optimal scan unroll is opposite on cpu and gpu (benchmarked):
222 # - gpu: the loop is dominated by per-step overhead, so a large unroll is up
223 # to ~6x faster; the run time plateaus by ~32 while compile time then grows
224 # steeply, so 32 is the sweet spot.
225 # - cpu: past ~6 the backend stops aliasing `out` in place and copies the
226 # size-`size` buffer each step (O(size**2), ~100x slower), so 2 is safest.
227 # `default` (cpu, tpu, untested backends) takes the conservative value.
228 return lax.platform_dependent(
229 cuda=partial(run, 32), rocm=partial(run, 32), default=partial(run, 2)
230 )
233class split:
234 """
235 Split a key into `num` keys.
237 Parameters
238 ----------
239 key
240 The key to split.
241 num
242 The number of keys to split into.
243 """
245 _keys: tuple[Key[Array, ''], ...]
246 _num_used: int
248 def __init__(self, key: Key[Array, ''], num: int = 2) -> None:
249 self._keys = _split_unpack(key, num)
250 self._num_used = 0
252 def __len__(self) -> int:
253 return len(self._keys) - self._num_used
255 def pop(self, shape: int | tuple[int, ...] = ()) -> Key[Array, ' *shape']:
256 """
257 Pop one or more keys from the list.
259 Parameters
260 ----------
261 shape
262 The shape of the keys to pop. If empty (default), a single key is
263 popped and returned. If not empty, the popped key is split and
264 reshaped to the target shape.
266 Returns
267 -------
268 The popped keys as a jax array with the requested shape.
270 Raises
271 ------
272 IndexError
273 If the list is empty.
274 """
275 if len(self) == 0:
276 msg = 'No keys left to pop'
277 raise IndexError(msg)
278 if not isinstance(shape, tuple):
279 shape = (shape,)
280 key = self._keys[self._num_used]
281 self._num_used += 1
282 if shape:
283 key = _split_shaped(key, shape)
284 return key
287@jit(static_argnums=(1,))
288def _split_unpack(key: Key[Array, ''], num: int) -> tuple[Key[Array, ''], ...]:
289 keys = random.split(key, num)
290 return tuple(keys)
293@jit(static_argnums=(1,))
294def _split_shaped(key: Key[Array, ''], shape: tuple[int, ...]) -> Key[Array, ' *shape']:
295 num = math.prod(shape)
296 keys = random.split(key, num)
297 return keys.reshape(shape)
300def truncated_normal_onesided(
301 key: Key[Array, ''],
302 shape: Sequence[int],
303 upper: Bool[Array, '...'],
304 bound: Float32[Array, '...'],
305 *,
306 clip: bool = True,
307) -> Float32[Array, '...']:
308 """
309 Sample from a one-sided truncated standard normal distribution.
311 Parameters
312 ----------
313 key
314 JAX random key.
315 shape
316 Shape of output array, broadcasted with other inputs.
317 upper
318 True for (-∞, bound], False for [bound, ∞).
319 bound
320 The truncation boundary.
321 clip
322 Whether to clip the samples to keep them finite and within the
323 truncation region. Leave on, off only for debugging.
325 Returns
326 -------
327 Array of samples from the truncated normal distribution.
328 """
329 # Pseudocode:
330 # | if upper:
331 # | if bound < 0:
332 # | ndtri(uniform(0, ndtr(bound))) =
333 # | ndtri(ndtr(bound) * u)
334 # | if bound > 0:
335 # | -ndtri(uniform(ndtr(-bound), 1)) =
336 # | -ndtri(ndtr(-bound) + ndtr(bound) * (1 - u))
337 # | if not upper:
338 # | if bound < 0:
339 # | ndtri(uniform(ndtr(bound), 1)) =
340 # | ndtri(ndtr(bound) + ndtr(-bound) * (1 - u))
341 # | if bound > 0:
342 # | -ndtri(uniform(0, ndtr(-bound))) =
343 # | -ndtri(ndtr(-bound) * u)
344 shape = jnp.broadcast_shapes(shape, upper.shape, bound.shape)
345 bound_pos = bound > 0
346 ndtr_bound = ndtr(bound)
347 ndtr_neg_bound = ndtr(-bound)
348 scale = jnp.where(upper, ndtr_bound, ndtr_neg_bound)
349 shift = jnp.where(upper, ndtr_neg_bound, ndtr_bound)
350 u = random.uniform(key, shape)
351 left_u = scale * (1 - u) # ~ uniform in (0, ndtr(±bound)]
352 right_u = shift + scale * u # ~ uniform in [ndtr(∓bound), 1)
353 truncated_u = jnp.where(upper ^ bound_pos, left_u, right_u)
354 if clip:
355 # `truncated_u` can reach 0 or 1, where ndtri is infinite: ndtr(±bound)
356 # underflows for extreme bounds, u can come out exactly 0, and on gpu the
357 # accuracy is lower. The lower target is the smallest normal number rather
358 # than the smallest subnormal one because xla flushes subnormals to zero on
359 # cpu, and ndtri is -inf on subnormals anyway.
360 zero = jnp.zeros((), truncated_u.dtype)
361 one = jnp.ones((), truncated_u.dtype)
362 smallest_normal = jnp.finfo(truncated_u.dtype).smallest_normal
363 truncated_u = jnp.clip(truncated_u, smallest_normal, jnp.nextafter(one, zero))
364 truncated_norm = ndtri(truncated_u)
365 sample = jnp.where(bound_pos, -truncated_norm, truncated_norm)
366 if clip:
367 # the clip above caps |sample| at ndtri of the target, ~12.9 in float32,
368 # which for extreme bounds falls short of the truncation region, and
369 # rounding can put the sample a ulp on the wrong side anyway. Move it onto
370 # the boundary, where the distribution concentrates for such bounds (its
371 # mean is bound + 1/bound).
372 sample = jnp.where(
373 upper, jnp.minimum(sample, bound), jnp.maximum(sample, bound)
374 )
375 return sample
378def get_default_device() -> Device:
379 """Get the current default JAX device."""
380 with ensure_compile_time_eval():
381 return jnp.empty(0).device
384def get_default_devices() -> list[Device]:
385 """Get all JAX devices on the default platform."""
386 return jax.devices(get_default_device().platform)
389def get_device_count() -> int:
390 """Get the number of available devices on the default platform."""
391 return len(get_default_devices())
394def is_key(x: object) -> TypeIs[Key[Array, ' *shape']]:
395 """Determine if `x` is a jax random key."""
396 return isinstance(x, Array) and jnp.issubdtype(x.dtype, prng_key)
399def jit_active() -> bool:
400 """Check if we are under jit."""
401 return not hasattr(jnp.empty(0), 'platform')
404def _equal_shards(x: Shaped[Array, '...'], axis_name: str) -> Bool[Array, '']:
405 """Check if all shards of `x` are equal, to be used in a `shard_map` context."""
406 size = lax.axis_size(axis_name)
407 perm = [(i, (i + 1) % size) for i in range(size)]
408 perm_x = lax.ppermute(x, axis_name, perm)
409 diff = jnp.any(x != perm_x)
410 return jnp.logical_not(lax.psum(diff, axis_name))
413def equal_shards(
414 x: PyTree[Array, ' S'], axis_name: str, **shard_map_kwargs: Any
415) -> PyTree[Bool[Array, ''], ' S']:
416 """Check that all shards of `x` are equal across axis `axis_name`.
418 Parameters
419 ----------
420 x
421 A pytree of arrays to check. Each array is checked separately.
422 axis_name
423 The mesh axis name across which equality is checked. It's not checked
424 across other axes.
425 **shard_map_kwargs
426 Additional arguments passed to `jax.shard_map` to set up the function
427 that checks equality. You may need to specify `in_specs` passing
428 the (pytree of) `jax.sharding.PartitionSpec` that specifies how `x`
429 is sharded, if the axes are not explicit, and `mesh` if there is not
430 a default mesh set by `jax.set_mesh`.
432 Returns
433 -------
434 A pytree of booleans indicating whether each leaf is equal across devices along the mesh axis.
435 """
436 equal_shards_leaf = partial(_equal_shards, axis_name=axis_name)
438 def check_equal(x: PyTree[Array, ' S']) -> PyTree[Bool[Array, ''], ' S']:
439 return tree.map(equal_shards_leaf, x)
441 sharded_check_equal = shard_map(
442 check_equal, out_specs=PartitionSpec(), **shard_map_kwargs
443 )
445 return sharded_check_equal(x)