Coverage for src/bartz/mcmcstep/_reduction.py: 99%

1# bartz/src/bartz/mcmcstep/_reduction.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. 

24 

25"""Indexed-reduce (scatter-add) configs, one per algorithm, and the core ops.""" 

26 

27import math (empty)

28from abc import abstractmethod (empty)

29from functools import partial (empty)

30from typing import Literal (empty)

31 

32import jax (empty)

33from equinox import Module, field (empty)

34from jax import lax (empty)

35from jax import numpy as jnp (empty)

36from jax.extend.backend import backends (empty)

37from jax.typing import DTypeLike (empty)

38from jaxtyping import Array, Float, Integer, Shaped, UInt (empty)

39 

40# target number of datapoint batches on cpu, and minimum datapoints per batch, 

41# when batching is resolved automatically; unlike the gpu heuristic these are 

42# flat (the cpu has no SM-count analog to scale with) 

43_AUTO_CPU_TARGET = 16 (empty)

44_AUTO_CPU_MIN_BATCH = 32 (empty)

45 

46# SM count used to trace `AutoBatchedReduction`'s gpu branch when no cuda backend 

47# is visible. `lax.platform_dependent` traces that branch even with no gpu, only 

48# to discard it at lowering, so this value never sizes a real gpu's batch grid; 

49# it just keeps the dead trace valid (a present cuda backend reports its own). 

50_MOOT_GPU_SM = 1 (empty)

51 

52 

53class ReductionConfig(Module): (empty)

54 """Select and configure an indexed-reduce (scatter-add) implementation. 

55 

56 Each concrete subclass identifies a reduction algorithm and carries its 

57 options. Pass instances to `init` to control how the residuals, counts and 

58 likelihood precisions are summed over the datapoints in each leaf. 

59 """ 

60 

61 @abstractmethod (empty)

62 def _reduce( (empty)

63 self, 

64 values: Float[Array, '*batch_shape n'] | int, 

65 indices: UInt[Array, ' n'], 

66 /, 

67 *, 

68 size: int, 

69 subset_start: Integer[Array, ''] | None = None, 

70 subset_length: int | None = None, 

71 dtype: DTypeLike, 

72 data_sharded: bool, 

73 # the output's trailing axis is the number of reduced bins: the range 

74 # length, or `size` without a subset. jaxtyping evals the `{...}` dim 

75 # against the arguments but forbids spaces and str-formats arrays, so 

76 # this indexes a tuple by a bool instead of `... if ... else ...`. The 

77 # bool reads a zero-length range as no subset, mislabeling the dim, but 

78 # the only caller passes the nonempty two-element child pair. 

79 ) -> Shaped[Array, '*batch_shape {(size,subset_length)[bool(subset_length)]}']: 

80 """Indexed reduce along the last axis of `values`. 

81 

82 Parameters 

83 ---------- 

84 values 

85 The values to sum into bins, or a scalar `int` weighting every 

86 datapoint equally (used to count the datapoints in each bin). 

87 indices 

88 The bin index each datapoint falls into. 

89 size 

90 The static excluded upper bound on the values of `indices`, and 

91 the number of output bins when no subset is given. 

92 subset_start 

93 If given (with `subset_length`), reduce only into the contiguous 

94 bin range ``[subset_start, subset_start + subset_length)``, 

95 ignoring datapoints whose index falls elsewhere. The range may run 

96 past `size`; those out-of-domain bins reduce to zero. 

97 subset_length 

98 The static length of the bin range, or `None` to reduce into all 

99 `size` bins. 

100 dtype 

101 The dtype of the output and of the accumulation; the values are 

102 kept in their own, possibly narrower, dtype until accumulated. 

103 data_sharded 

104 Whether the data axis is sharded; if true, the result is 

105 psum-reduced across the ``'data'`` axis of the enclosing 

106 `shard_map`. 

107 

108 Returns 

109 ------- 

110 The per-bin sums, with the same leading dimensions as `values` and the bins on the trailing axis. 

111 """ 

112 ... 

113 

114 

115def _resolve_range( (empty)

116 indices: UInt[Array, ' n'], 

117 size: int, 

118 subset_start: Integer[Array, ''] | None, 

119 subset_length: int | None, 

120) -> tuple[int, UInt[Array, ' n']]: 

121 """Reduce the contiguous-range subset to the full case for scatter algorithms. 

122 

123 Parameters 

124 ---------- 

125 indices 

126 The bin index each datapoint falls into, in ``[0, size)``. 

127 size 

128 The number of bins. 

129 subset_start 

130 The first bin of the range to reduce into, or `None` for all bins. 

131 subset_length 

132 The static number of bins in the range, or `None` for all bins. 

133 

134 Returns 

135 ------- 

136 out_size : int 

137 The number of output bins: `subset_length`, or `size` without a subset. 

138 indices : UInt[Array, ' n'] 

139 The scatter indices into the output bins: unchanged without a subset, 

140 else each datapoint's offset from `subset_start`, in the indices' own 

141 unsigned dtype, so that indices outside ``[subset_start, subset_start + 

142 subset_length)`` land out of bounds, where the scatter drops them. 

143 """ 

144 if subset_length is None: (empty)

145 return size, indices (empty)

146 else: 

147 # the subtraction is unsigned: indices below `subset_start` underflow to 

148 # a large value rather than going negative (which the scatter would read 

149 # as wrap-around indexing), and together with indices ``>= subset_start + 

150 # subset_length`` they fall outside the output, where scatters drop them. 

151 # Exact while the range fits the index dtype (``subset_start + 

152 # subset_length <= 2 ** bits``), as it does for the per-move child pair. 

153 assert subset_start is not None # set together with subset_length (empty)

154 assert jnp.issubdtype(indices.dtype, jnp.unsignedinteger) (empty)

155 offset = indices - subset_start.astype(indices.dtype) (empty)

156 return subset_length, offset (empty)

157 

158 

159class BatchedReduction(ReductionConfig): (empty)

160 """Segment-sum with optional batching along the datapoints. 

161 

162 Fastest at the usual tree sizes. See `AutoBatchedReduction` to resolve 

163 `num_batches` automatically per platform. 

164 """ 

165 

166 num_batches: int | None = field(static=True, default=None) (empty)

167 """The number of datapoint batches, or `None` (the default) for an unbatched (empty)

168 reduce.""" 

169 

170 batches_inner: bool = field(static=True, default=True) (empty)

171 """Whether the batch axis sits on the scatter buffer's inner, contiguous axis (empty)

172 (``size``-by-``num_batches``) or its outer axis (``num_batches``-by-``size``); 

173 the two layouts give the backend different memory access patterns. `True` (the 

174 default) matches the historical layout. No effect when `num_batches` is `None`.""" 

175 

176 contiguous: bool = field(static=True, default=False) (empty)

177 """How datapoints are assigned to batches. `False` (the default) strides them, (empty)

178 sending datapoint ``i`` to batch ``i % num_batches``; `True` splits them into 

179 contiguous chunks, sending ``i`` to batch ``i // batch_size``. No effect when 

180 `num_batches` is `None`.""" 

181 

182 def _reduce( (empty)

183 self, 

184 values: Float[Array, '*batch_shape n'] | int, 

185 indices: UInt[Array, ' n'], 

186 /, 

187 *, 

188 size: int, 

189 subset_start: Integer[Array, ''] | None = None, 

190 subset_length: int | None = None, 

191 dtype: DTypeLike, 

192 data_sharded: bool, 

193 ) -> Shaped[Array, '*batch_shape {(size,subset_length)[bool(subset_length)]}']: 

194 values = jnp.asarray(values) (empty)

195 assert values.ndim == 0 or values.shape[-1:] == indices.shape (empty)

196 size, indices = _resolve_range(indices, size, subset_start, subset_length) (empty)

197 batch_shape = values.shape[:-1] (empty)

198 

199 if self.num_batches is None: (empty)

200 out = jnp.zeros((*batch_shape, size), dtype).at[..., indices].add(values) (empty)

201 else: 

202 # in the sharded case, n is the size of the local shard, not the full size 

203 (n,) = indices.shape (empty)

204 # unsigned avoids a negative-index normalization select in the scatter 

205 iota = jnp.arange(n, dtype=jnp.uint32) (empty)

206 if self.contiguous: (empty)

207 batch_size = -(-n // self.num_batches) # ceil: last batch is partial (empty)

208 batch_indices = iota // batch_size (empty)

209 else: 

210 batch_indices = iota % self.num_batches (empty)

211 if self.batches_inner: (empty)

212 out = ( (empty)

213 jnp.zeros((*batch_shape, size, self.num_batches), dtype) 

214 .at[..., indices, batch_indices] 

215 .add(values) 

216 .sum(axis=-1) 

217 ) 

218 else: 

219 out = ( (empty)

220 jnp.zeros((*batch_shape, self.num_batches, size), dtype) 

221 .at[..., batch_indices, indices] 

222 .add(values) 

223 .sum(axis=-2) 

224 ) 

225 

226 if data_sharded: (empty)

227 out = lax.psum(out, 'data') (empty)

228 return out (empty)

229 

230 

231class AutoBatchedReduction(ReductionConfig): (empty)

232 """`BatchedReduction` that picks `num_batches` automatically per platform. 

233 

234 A flat target on cpu, and on gpu a count scaling with the SM count and the 

235 multivariate outcome size. Only cpu and cuda are supported; any other 

236 platform raises at lowering. 

237 """ 

238 

239 min_batch_size: float = field(static=True, default=128.0) (empty)

240 """Minimum datapoints per batch on gpu; caps the batch count at (empty)

241 ``n / min_batch_size``.""" 

242 

243 beta_sm: float = field(static=True, default=48.0) (empty)

244 """Batches per streaming multiprocessor on gpu; the batch count saturates at (empty)

245 ``beta_sm * n_sms * m ** -gamma``, with `n_sms` the gpu's SM count and ``m`` 

246 the multivariate work per datapoint.""" 

247 

248 gamma: float = field(static=True, default=0.4) (empty)

249 """Exponent by which multivariate outcomes (``m`` values per datapoint) shrink (empty)

250 the saturation batch count on gpu.""" 

251 

252 batches_inner: bool = field(static=True, default=True) (empty)

253 """Same as `BatchedReduction`.""" (empty)

254 

255 contiguous: bool = field(static=True, default=False) (empty)

256 """Same as `BatchedReduction`.""" (empty)

257 

258 def _reduce( (empty)

259 self, 

260 values: Float[Array, '*batch_shape n'] | int, 

261 indices: UInt[Array, ' n'], 

262 /, 

263 *, 

264 size: int, 

265 subset_start: Integer[Array, ''] | None = None, 

266 subset_length: int | None = None, 

267 dtype: DTypeLike, 

268 data_sharded: bool, 

269 ) -> Shaped[Array, '*batch_shape {(size,subset_length)[bool(subset_length)]}']: 

270 # defer the cpu/gpu choice to XLA: both branches are traced, but only the 

271 # one for the run platform is lowered. With no `default`, an untested 

272 # platform (rocm, tpu) errors at lowering instead of silently falling back. 

273 kwargs = dict( (empty)

274 size=size, 

275 subset_start=subset_start, 

276 subset_length=subset_length, 

277 dtype=dtype, 

278 data_sharded=data_sharded, 

279 ) 

280 return lax.platform_dependent( (empty)

281 cpu=partial(self._reduce_cpu, values, indices, **kwargs), 

282 cuda=partial(self._reduce_gpu, values, indices, **kwargs), 

283 ) 

284 

285 def _reduce_cpu( (empty)

286 self, 

287 values: Float[Array, '*batch_shape n'] | int, 

288 indices: UInt[Array, ' n'], 

289 /, 

290 *, 

291 size: int, 

292 subset_start: Integer[Array, ''] | None = None, 

293 subset_length: int | None = None, 

294 dtype: DTypeLike, 

295 data_sharded: bool, 

296 ) -> Shaped[Array, '*batch_shape {(size,subset_length)[bool(subset_length)]}']: 

297 # flat target: the cpu has no SM-count analog to scale the batch count with 

298 (n,) = indices.shape (empty)

299 num_batches = _final_round(n, _AUTO_CPU_TARGET, _AUTO_CPU_MIN_BATCH) (empty)

300 return self._delegate(num_batches)._reduce( # noqa: SLF001 (empty)

301 values, 

302 indices, 

303 size=size, 

304 subset_start=subset_start, 

305 subset_length=subset_length, 

306 dtype=dtype, 

307 data_sharded=data_sharded, 

308 ) 

309 

310 def _reduce_gpu( (empty)

311 self, 

312 values: Float[Array, '*batch_shape n'] | int, 

313 indices: UInt[Array, ' n'], 

314 /, 

315 *, 

316 size: int, 

317 subset_start: Integer[Array, ''] | None = None, 

318 subset_length: int | None = None, 

319 dtype: DTypeLike, 

320 data_sharded: bool, 

321 ) -> Shaped[Array, '*batch_shape {(size,subset_length)[bool(subset_length)]}']: 

322 # `n` is the local shard size when data-sharded, which is exactly what the 

323 # heuristic wants. `m` is the multivariate batch size: the product of the 

324 # values' leading (non-datapoint) axes (1 for the scalar count case), i.e. 

325 # how much vector work rides on each scatter slot; clamped to >=1 so a 

326 # zero-size leading axis (e.g. k=0 outcome components) leaves a moot, empty 

327 # reduce with the m=1 baseline cap. 

328 (n,) = indices.shape (empty)

329 m = max(1, math.prod(jnp.shape(values)[:-1])) (empty)

330 # the gpu cap scales up with the SM count and down, sublinearly, with the 

331 # multivariate work per slot 

332 sm_cap = self.beta_sm * _gpu_sm_count() * m ** (-self.gamma) (empty)

333 num_batches = _final_round(n, sm_cap, self.min_batch_size) (empty)

334 return self._delegate(num_batches)._reduce( # noqa: SLF001 (empty)

335 values, 

336 indices, 

337 size=size, 

338 subset_start=subset_start, 

339 subset_length=subset_length, 

340 dtype=dtype, 

341 data_sharded=data_sharded, 

342 ) 

343 

344 def _delegate(self, num_batches: int | None) -> BatchedReduction: (empty)

345 """Build a `BatchedReduction` with the resolved count and this config's layout.""" 

346 return BatchedReduction( (empty)

347 num_batches=num_batches, 

348 batches_inner=self.batches_inner, 

349 contiguous=self.contiguous, 

350 ) 

351 

352 

353def _final_round( (empty)

354 n: int, target: float | int, min_batch_size: float | int 

355) -> int | None: 

356 """Cap batches to keep them above `min_batch_size`, round to a power of 2, and disable batching if there's only 1 batch.""" 

357 # at least `min_batch_size` elements per batch 

358 num = min(n / min_batch_size, target) (empty)

359 

360 # round to the nearest power of 2 because I guess XLA and the hardware 

361 # will like that (not sure about this, maybe just multiple of 32?) 

362 num = 2 ** round(math.log2(num)) if num > 0 else 0 (empty)

363 

364 # disable batching if the batch is as large as the whole dataset 

365 return num if num > 1 else None (empty)

366 

367 

368def _gpu_sm_count() -> int: (empty)

369 """Streaming-multiprocessor count shared by the visible cuda gpus. 

370 

371 Read by `AutoBatchedReduction` to size the gpu batch grid. Since 

372 `lax.platform_dependent` only runs the gpu branch on cuda, this trusts 

373 `jax.devices('cuda')` and each device's `core_count` rather than guessing, 

374 and raises if the gpus report differing counts (a mixed gpu set is 

375 unsupported). 

376 """ 

377 if 'cuda' not in backends(): (empty)

378 # no cuda backend: lax.platform_dependent still traces the gpu branch 

379 # here, only to discard it at lowering, so the count is never used 

380 return _MOOT_GPU_SM (empty)

381 counts = {device.core_count for device in jax.devices('cuda')} (empty)

382 if len(counts) > 1: (empty)

383 msg = ( (empty)

384 f'visible cuda gpus report differing SM counts {sorted(counts)}; ' 

385 'AutoBatchedReduction assumes a single gpu model' 

386 ) 

387 raise ValueError(msg) (empty)

388 (count,) = counts (empty)

389 return count (empty)

390 

391 

392class OneHotReduction(ReductionConfig): (empty)

393 """Dense one-hot reduction. 

394 

395 Materializes the membership of each datapoint in its leaf as a one-hot 

396 matrix over the output bins and contracts it against the values. Beats 

397 `BatchedReduction` only when the number of bins is very small (e.g. a 

398 single leaf pair), or on gpu for multivariate residuals. 

399 """ 

400 

401 method: Literal['matmul', 'multiply', 'scatter_set'] = field( (empty)

402 static=True, default='matmul' 

403 ) 

404 """How to contract the values against the one-hot leaf-membership matrix: (empty)

405 

406 'matmul' 

407 Contract the values with the one-hot matrix via a dot. Faster on gpu, 

408 especially for multivariate residuals. 

409 'multiply' 

410 Elementwise-multiply by the one-hot matrix and reduce over the 

411 datapoints; whether the ``n``-by-``size`` product is fused into the 

412 reduction or materialized is left to the backend. Faster on cpu. 

413 'scatter_set' 

414 Scatter the values into a dense buffer with unique (non-atomic) writes, 

415 then sum over the datapoints. 

416 """ 

417 

418 n_inner: bool = field(static=True, default=True) (empty)

419 """Whether the datapoints sit on the one-hot's inner, contiguous axis (empty)

420 (``size``-by-``n``) or its outer axis (``n``-by-``size``); the two layouts 

421 give the backend different memory access patterns. `True` (the default) 

422 fuses better on gpu.""" 

423 

424 def _reduce( (empty)

425 self, 

426 values: Float[Array, '*batch_shape n'] | int, 

427 indices: UInt[Array, ' n'], 

428 /, 

429 *, 

430 size: int, 

431 subset_start: Integer[Array, ''] | None = None, 

432 subset_length: int | None = None, 

433 dtype: DTypeLike, 

434 data_sharded: bool, 

435 ) -> Shaped[Array, '*batch_shape {(size,subset_length)[bool(subset_length)]}']: 

436 values = jnp.asarray(values) (empty)

437 assert values.ndim == 0 or values.shape[-1:] == indices.shape (empty)

438 

439 # a scalar value is the count case, weighting each datapoint by `values`; 

440 # it broadcasts in the scatter/multiply paths, only matmul needs a vector 

441 scalar = values.ndim == 0 (empty)

442 (n,) = indices.shape (empty)

443 batch_shape = values.shape[:-1] (empty)

444 size, bins, indices = _resolve_range_bins( (empty)

445 indices, 

446 size, 

447 subset_start, 

448 subset_length, 

449 remap=self.method == 'scatter_set', 

450 ) 

451 # unsigned avoids a negative-index normalization select in the scatter 

452 iota = jnp.arange(n, dtype=jnp.uint32) (empty)

453 

454 # one-hots and scatter buffers hold the values, so they are built in 

455 # the values' dtype; only the reduction accumulates in `dtype`. The 

456 # scalar count case has no input precision to preserve and uses `dtype`. 

457 values_dtype = dtype if scalar else values.dtype (empty)

458 

459 match self.method, self.n_inner: (empty)

460 case 'scatter_set', True: (empty)

461 out = ( (empty)

462 jnp.zeros((*batch_shape, size, n), values_dtype) 

463 .at[..., indices, iota] 

464 .set(values, unique_indices=True) 

465 .sum(axis=-1, dtype=dtype) 

466 ) 

467 case 'scatter_set', False: (empty)

468 out = ( (empty)

469 jnp.zeros((*batch_shape, n, size), values_dtype) 

470 .at[..., iota, indices] 

471 .set(values, unique_indices=True) 

472 .sum(axis=-2, dtype=dtype) 

473 ) 

474 case 'matmul', True: (empty)

475 onehot = (bins[:, None] == indices).astype(values_dtype) # (size, n) (empty)

476 vec = jnp.broadcast_to(values.astype(dtype), (n,)) if scalar else values (empty)

477 out = jnp.einsum( (empty)

478 '...n,sn->...s', vec, onehot, preferred_element_type=dtype 

479 ) 

480 case 'matmul', False: (empty)

481 onehot = (indices[:, None] == bins).astype(values_dtype) # (n, size) (empty)

482 vec = jnp.broadcast_to(values.astype(dtype), (n,)) if scalar else values (empty)

483 out = jnp.einsum( (empty)

484 '...n,ns->...s', vec, onehot, preferred_element_type=dtype 

485 ) 

486 case 'multiply', True: (empty)

487 onehot = bins[:, None] == indices # (size, n) (empty)

488 if scalar: (empty)

489 out = values * onehot.sum(axis=-1, dtype=dtype) (empty)

490 else: 

491 out = (values[..., None, :] * onehot).sum(axis=-1, dtype=dtype) (empty)

492 case 'multiply', False: 492 ↛ 499line 492 didn't jump to line 499 because the pattern on line 492 always matched(empty)

493 onehot = indices[:, None] == bins # (n, size) (empty)

494 if scalar: (empty)

495 out = values * onehot.sum(axis=-2, dtype=dtype) (empty)

496 else: 

497 out = (values[..., :, None] * onehot).sum(axis=-2, dtype=dtype) (empty)

498 

499 if data_sharded: (empty)

500 out = lax.psum(out, 'data') (empty)

501 return out (empty)

502 

503 

504def _resolve_range_bins( (empty)

505 indices: UInt[Array, ' n'], 

506 size: int, 

507 subset_start: Integer[Array, ''] | None, 

508 subset_length: int | None, 

509 *, 

510 remap: bool, 

511) -> tuple[int, UInt[Array, ' out_size'], UInt[Array, ' n']]: 

512 """Resolve the range subset into output size, comparison bins, and scatter indices. 

513 

514 The comparison methods of `OneHotReduction` reduce against the range's bins 

515 directly, while its scatter method (`remap`) indexes bins by position, so 

516 the indices are offset like in `_resolve_range`. 

517 """ 

518 if subset_length is None: (empty)

519 return size, jnp.arange(size, dtype=indices.dtype), indices (empty)

520 assert subset_start is not None # set together with subset_length (empty)

521 # uint32, not the possibly narrow `indices.dtype`, so bins past `size` do 

522 # not wrap and alias a real bin in the comparison 

523 bins = subset_start.astype(jnp.uint32) + jnp.arange(subset_length, dtype=jnp.uint32) (empty)

524 if remap: (empty)

525 out_size, indices = _resolve_range(indices, size, subset_start, subset_length) (empty)

526 return out_size, bins, indices (empty)

527 else: 

528 return subset_length, bins, indices (empty)

529 

530 

531class AutoOneHotReduction(ReductionConfig): (empty)

532 """`OneHotReduction` that picks `method` and `n_inner` automatically. 

533 

534 Resolves both knobs from trace-time information per site and platform, then 

535 delegates to a plain `OneHotReduction`. Uses `matmul` only for wide-bin 

536 multivariate reductions and `multiply` otherwise; lays the datapoints on the 

537 outer axis except on the two small-bin sites where the opposite wins (cpu 

538 precision, cuda count). Those two sites support only cpu and cuda, raising at 

539 lowering elsewhere. 

540 

541 The site is recovered from the value: scalar is the count, a wide output the 

542 residual, a narrow non-scalar output the precision. 

543 

544 Known limitation: the wide-bin univariate residual on cpu past ~10^6 

545 datapoints prefers a layout this picks against (up to ~2x slower). 

546 """ 

547 

548 min_matmul_bins: int = field(static=True, default=8) (empty)

549 """Minimum output bins for `matmul`; below it `multiply` is always used.""" (empty)

550 

551 def _reduce( (empty)

552 self, 

553 values: Float[Array, '*batch_shape n'] | int, 

554 indices: UInt[Array, ' n'], 

555 /, 

556 *, 

557 size: int, 

558 subset_start: Integer[Array, ''] | None = None, 

559 subset_length: int | None = None, 

560 dtype: DTypeLike, 

561 data_sharded: bool, 

562 ) -> Shaped[Array, '*batch_shape {(size,subset_length)[bool(subset_length)]}']: 

563 out_size = size if subset_length is None else subset_length (empty)

564 m = max(1, math.prod(jnp.shape(values)[:-1])) (empty)

565 method = 'matmul' if m >= 2 and out_size >= self.min_matmul_bins else 'multiply' (empty)

566 

567 if jnp.ndim(values) == 0: # count (empty)

568 cpu_inner, cuda_inner = False, True (empty)

569 elif out_size <= 2: # precision (empty)

570 cpu_inner, cuda_inner = True, False (empty)

571 else: # residual 

572 cpu_inner, cuda_inner = False, False (empty)

573 

574 args = (values, indices) (empty)

575 kwargs: dict = dict( (empty)

576 size=size, 

577 subset_start=subset_start, 

578 subset_length=subset_length, 

579 dtype=dtype, 

580 data_sharded=data_sharded, 

581 ) 

582 if cpu_inner == cuda_inner: (empty)

583 # the layout matches on every platform, so no platform split is 

584 # needed and the reduction also runs on untested platforms (tpu/rocm) 

585 return OneHotReduction(method=method, n_inner=cpu_inner)._reduce( # noqa: SLF001 (empty)

586 *args, **kwargs 

587 ) 

588 else: 

589 # defer the cpu/gpu choice to XLA: both branches are traced, but only 

590 # the run platform's is lowered. With no `default`, an untested 

591 # platform errors at lowering instead of silently falling back. 

592 return lax.platform_dependent( (empty)

593 cpu=partial( 

594 OneHotReduction(method=method, n_inner=cpu_inner)._reduce, # noqa: SLF001 

595 *args, 

596 **kwargs, 

597 ), 

598 cuda=partial( 

599 OneHotReduction(method=method, n_inner=cuda_inner)._reduce, # noqa: SLF001 

600 *args, 

601 **kwargs, 

602 ), 

603 )