mlboydaisuke
Timestral-209M on LiteRT: float16 graphs for the GPU and the NPU, int8 graphs for the CPU, host code in Python and Kotlin
11f5cc3 verified 13 days ago
Raw History Blame Contribute Delete
26 kB
package com.timestral

import kotlin.math.max
import kotlin.math.min

/**
 * One exported graph shape: [contextPatches] context patches and [horizonPatches] decoder patches,
 * so the graph always sees `contextPatches + horizonPatches` tokens of 32 steps each.
 */
data class Bucket(val name: String, val contextPatches: Int, val horizonPatches: Int) {
  /** Tokens of the graph (S = L + H). */
  val tokens: Int
    get() = contextPatches + horizonPatches
}

/**
 * The host tables installed next to the graph files: the checkpoint's `forecast_query` vector (768
 * floats, the decoder token) and the RoPE cos / sin tables (rows of 64 floats, row 0 at position
 * [firstPosition]), written with the official `SinCosRoPE` module by the conversion run.
 */
class HostTables(
  val forecastQuery: FloatArray,
  val ropeCos: FloatArray,
  val ropeSin: FloatArray,
  val firstPosition: Int,
) {
  init {
    require(forecastQuery.size == TimestralHost.D_MODEL) {
      "forecast_query has ${forecastQuery.size} values, expected ${TimestralHost.D_MODEL}"
    }
    require(ropeCos.size == ropeSin.size && ropeCos.size % TimestralHost.ROPE_DIM == 0) {
      "RoPE tables have ${ropeCos.size} and ${ropeSin.size} values"
    }
  }

  /** Rows of each RoPE table. */
  val rows: Int
    get() = ropeCos.size / TimestralHost.ROPE_DIM
}

/**
 * The five float32 inputs of the graph's `serving_default` signature, flattened row-major: patches
 * `[R, L, 64]`, dec `[R, H, 768]`, bias `[1, 1, S, S]`, cos and sin `[1, 1, S, 64]`.
 */
class GraphInputs(
  val patches: FloatArray,
  val dec: FloatArray,
  val bias: FloatArray,
  val cos: FloatArray,
  val sin: FloatArray,
  /** The first context patch (of the bucket's L) with an observation; 0 when none has one. */
  val firstObservedPatch: Int,
) {
  /** The inputs by signature name. */
  fun byName(): Map<String, FloatArray> =
    mapOf("patches" to patches, "dec" to dec, "bias" to bias, "cos" to cos, "sin" to sin)
}

/**
 * One series made ready for the graph: the context row (missing steps forward-filled, the last
 * [TimestralHost.MAX_CONTEXT] steps, left-padded to whole patches), its negation (the mirror row of
 * symmetric averaging), the causal statistics of both rows and their normalized values.
 */
class PreparedSeries(
  /** Forecast steps asked for. */
  val horizon: Int,
  /** Decoder patches the forecast needs (h = ceil(horizon / 32)). */
  val horizonPatches: Int,
  /** The whole series with missing steps set to 0 (the non-negative check reads it). */
  val inputs: FloatArray,
  /** Which steps of the whole series are observed (finite). */
  val observedAll: BooleanArray,
  /** The context row: T steps, T a multiple of 32. */
  val context: FloatArray,
  /** Which steps of [context] are observed. */
  val observed: BooleanArray,
  /** Per row (context, mirror) and step: the causal location of the step's patch. */
  val loc: Array<FloatArray>,
  /** Per row and step: the causal scale of the step's patch. */
  val scale: Array<FloatArray>,
  /** Per row and step: asinh((x - loc) / scale) on observed steps, 0 elsewhere. */
  val scaled: Array<FloatArray>,
) {
  /** Context patches (T / 32). */
  val contextPatches: Int
    get() = context.size / TimestralHost.PATCH

  /** The row values: the context, then its negation. */
  fun row(index: Int): FloatArray = if (index == 0) context else negate(context)

  private fun negate(values: FloatArray) = FloatArray(values.size) { -values[it] }
}

/** The host refused a forecast (the official forecaster raises the same error). */
class ForecastException(message: String) : Exception(message)

/**
 * The host side of Timestral-209M around one graph call, ported from the conversion run's
 * `scripts/host_forecast.py` with the official forecaster's defaults: symmetric averaging (the
 * context and its negation as the two rows of one call), non-negative clipping, quantile cap k =
 * 1e4, one decoder block of up to 32 patches with the advancing refine frame. Each stage keeps the
 * official dtype: the causal statistics and the refine frames are float64, normalization,
 * denormalization, the cap and the averaging are float32, asinh and sinh are evaluated in float64
 * and rounded to float32. Sums of 32 values follow NumPy's pairwise order and `max`, `min` and
 * `clip` follow NumPy's rules for NaN and signed zeros, so that the JVM tests can hold the results
 * to the Python reference value by value.
 *
 * Android-free: the JVM tests run it against vectors the Python reference wrote.
 */
object TimestralHost {
  /** Steps per patch. */
  const val PATCH = 32

  /** Width of the model and of the decoder token. */
  const val D_MODEL = 768

  /** Columns of each RoPE table row. */
  const val ROPE_DIM = 64

  /** Quantile levels 0.01, 0.05, 0.1, 0.15 … 0.9, 0.95, 0.99. */
  const val QUANTILES = 21

  /** Index of the 0.5 level. */
  const val MEDIAN = 10

  /** Index of the 0.1 level. */
  const val P10 = 2

  /** Index of the 0.9 level. */
  const val P90 = 18

  /** Rows of every graph call: the context and its negation. */
  const val ROWS = 2

  /** Decoder patches of one block. */
  const val BLOCK_PATCHES = 32

  /** Longest horizon of one block: 1,024 steps. */
  const val MAX_HORIZON = BLOCK_PATCHES * PATCH

  /**
   * Longest context: the last 8,192 steps (256 patches, the largest graph of this sample), the same
   * as the official forecaster with `max_context=8192`.
   */
  const val MAX_CONTEXT = 8192

  /** Floor of the causal standard deviation (float64). */
  private const val MIN_SCALE = 1e-6

  /** Quantile cap: context range ± k × the last scale. */
  private const val CAP_K = 1e4f

  private const val F32_MAX = Float.MAX_VALUE

  /** float32 machine epsilon, 2^-23. */
  private val F32_EPS = Math.ulp(1f)

  /** Smallest positive normal float64. */
  private const val F64_TINY = java.lang.Double.MIN_NORMAL

  /** The refine frame's clip, `float32 max ** 0.5 / 1e3` in Python (1.8446743e16). */
  private val FOLD_CLIP = Math.sqrt(F32_MAX.toDouble()) / 1e3

  /**
   * Builds the rows of one forecast from [series] (NaN or infinite = missing) for [horizon] steps
   * (1 to [MAX_HORIZON]).
   */
  fun prepare(series: FloatArray, horizon: Int): PreparedSeries {
    require(horizon in 1..MAX_HORIZON) { "horizon $horizon is outside 1..$MAX_HORIZON" }
    require(series.isNotEmpty()) { "empty series" }
    val observedAll = BooleanArray(series.size) { series[it].isFinite() }
    val inputs = FloatArray(series.size) { if (observedAll[it]) series[it] else 0f }
    var rows = forwardFill(inputs, observedAll)
    var observed = observedAll
    if (rows.size > MAX_CONTEXT) {
      rows = rows.copyOfRange(rows.size - MAX_CONTEXT, rows.size)
      observed = observed.copyOfRange(observed.size - MAX_CONTEXT, observed.size)
    }
    val padded = padToPatch(rows, observed)
    val context = padded.first
    val contextObserved = padded.second
    val both = arrayOf(context, FloatArray(context.size) { -context[it] })
    val loc = arrayOfNulls<FloatArray>(ROWS)
    val scale = arrayOfNulls<FloatArray>(ROWS)
    val scaled = arrayOfNulls<FloatArray>(ROWS)
    for (row in 0 until ROWS) {
      val stats = patchLocScale(both[row], contextObserved)
      loc[row] = stats.first
      scale[row] = stats.second
      scaled[row] = normalize(both[row], stats.first, stats.second, contextObserved)
    }
    return PreparedSeries(
      horizon,
      min(BLOCK_PATCHES, (horizon + PATCH - 1) / PATCH),
      inputs,
      observedAll,
      context,
      contextObserved,
      loc.requireNoNulls(),
      scale.requireNoNulls(),
      scaled.requireNoNulls(),
    )
  }

  /**
   * The graph inputs of [prepared] for [bucket] and the series rows [rows] (0 = the context, 1 =
   * its negation; both for a two-row graph, one per call for a one-row graph): the context patches
   * right-aligned in the L slots (the slots in front are padding: value 0, every step missing), the
   * decoder tokens, the additive attention mask (0 where a pair may attend, [mask] where it may
   * not) and the RoPE rows of the positions counted from the first observed patch. The mask and
   * the RoPE rows are the same for both rows.
   */
  fun buildInputs(
    prepared: PreparedSeries,
    bucket: Bucket,
    tables: HostTables,
    mask: Float,
    rows: IntArray = intArrayOf(0, 1),
  ): GraphInputs {
    val l = bucket.contextPatches
    val h = bucket.horizonPatches
    val lc = prepared.contextPatches
    require(lc <= l && prepared.horizonPatches <= h) { "the series does not fit ${bucket.name}" }
    require(rows.isNotEmpty() && rows.all { it in 0 until ROWS }) { "rows ${rows.toList()}" }
    val pad = l - lc
    val s = l + h
    val width = 2 * PATCH
    val patches = FloatArray(rows.size * l * width)
    for ((slot, row) in rows.withIndex()) {
      val base = slot * l * width
      for (p in 0 until pad) {
        for (j in 0 until PATCH) patches[base + p * width + PATCH + j] = 1f
      }
      val scaled = prepared.scaled[row]
      for (p in 0 until lc) {
        val slot = base + (pad + p) * width
        for (j in 0 until PATCH) {
          val t = p * PATCH + j
          patches[slot + j] = scaled[t]
          patches[slot + PATCH + j] = if (prepared.observed[t]) 0f else 1f
        }
      }
    }
    val observedPatch = BooleanArray(l)
    for (p in 0 until lc) {
      var any = false
      for (j in 0 until PATCH) any = any || prepared.observed[p * PATCH + j]
      observedPatch[pad + p] = any
    }
    // A key is valid once any patch at or before it has an observation.
    val keyValid = BooleanArray(l)
    var seen = false
    for (p in 0 until l) {
      seen = seen || observedPatch[p]
      keyValid[p] = seen
    }
    val firstObserved = observedPatch.indexOfFirst { it }.let { if (it < 0) 0 else it }
    val bias = FloatArray(s * s)
    for (q in 0 until s) {
      for (k in 0 until s) {
        val allow =
          when {
            q < l && k < l -> (k <= q && keyValid[k]) || k == q
            q >= l && k < l -> keyValid[k]
            q >= l -> k <= q
            else -> false
          }
        bias[q * s + k] = if (allow) 0f else mask
      }
    }
    val cos = FloatArray(s * ROPE_DIM)
    val sin = FloatArray(s * ROPE_DIM)
    for (t in 0 until s) {
      val tableRow = t - firstObserved - tables.firstPosition
      require(tableRow in 0 until tables.rows) { "RoPE table too small for position $tableRow" }
      System.arraycopy(tables.ropeCos, tableRow * ROPE_DIM, cos, t * ROPE_DIM, ROPE_DIM)
      System.arraycopy(tables.ropeSin, tableRow * ROPE_DIM, sin, t * ROPE_DIM, ROPE_DIM)
    }
    val dec = FloatArray(rows.size * h * D_MODEL)
    for (token in 0 until rows.size * h) {
      System.arraycopy(tables.forecastQuery, 0, dec, token * D_MODEL, D_MODEL)
    }
    return GraphInputs(patches, dec, bias, cos, sin, firstObserved)
  }

  /**
   * The forecast from the graph output [raw] (`[2, H, 32, 21]` normalized quantiles of [bucket],
   * the context's row then the negation's; a one-row graph's two outputs one after the other):
   * denormalize, refine, cap, average the two rows and clip; `[horizon × 21]` row-major, step
   * after step. Throws [ForecastException] when a quantile is NaN (the context's row is checked
   * first, as the official forecaster does).
   */
  fun finish(prepared: PreparedSeries, raw: FloatArray, bucket: Bucket): FloatArray {
    val h = prepared.horizonPatches
    val perRow = bucket.horizonPatches * PATCH * QUANTILES
    require(raw.size == ROWS * perRow) { "graph output has ${raw.size} values, expected ${ROWS * perRow}" }
    val steps = h * PATCH
    val last = prepared.context.size - 1
    val rows = Array(ROWS) { row ->
      val values = FloatArray(steps * QUANTILES)
      for (i in values.indices) {
        values[i] = clip(raw[row * perRow + i], -F32_MAX, F32_MAX)
      }
      var q = denormalize(values, prepared.loc[row][last], prepared.scale[row][last])
      q = refineBlock(q, prepared.row(row), prepared.observed, h)
      q.copyOf(prepared.horizon * QUANTILES)
    }
    val forward = cap(rows[0], prepared.row(0), prepared.observed, prepared.scale[0][last])
    checkNoNan(forward)
    val mirror = cap(rows[1], prepared.row(1), prepared.observed, prepared.scale[1][last])
    checkNoNan(mirror)
    val out = FloatArray(forward.size)
    for (t in 0 until prepared.horizon) {
      for (i in 0 until QUANTILES) {
        val a = forward[t * QUANTILES + i]
        val b = -mirror[t * QUANTILES + (QUANTILES - 1 - i)]
        out[t * QUANTILES + i] = 0.5f * (a + b)
      }
    }
    if (nonNegativeContext(prepared.inputs, prepared.observedAll)) {
      for (i in out.indices) out[i] = npMax(out[i], 0f)
    }
    return out
  }

  // ---------------------------------------------------------------- context

  /** Missing steps take the last observed value before them (0 before the first one). */
  internal fun forwardFill(inputs: FloatArray, observed: BooleanArray): FloatArray {
    if (observed.all { it }) return inputs
    var source = 0
    return FloatArray(inputs.size) { t ->
      if (observed[t]) source = t
      inputs[source]
    }
  }

  /** Left-pads with missing zeros to a whole number of patches. */
  internal fun padToPatch(
    values: FloatArray,
    observed: BooleanArray,
  ): Pair<FloatArray, BooleanArray> {
    val extra = values.size % PATCH
    if (extra == 0) return values to observed
    val pad = PATCH - extra
    val padded = FloatArray(pad + values.size)
    System.arraycopy(values, 0, padded, pad, values.size)
    val mask = BooleanArray(pad + observed.size)
    System.arraycopy(observed, 0, mask, pad, observed.size)
    return padded to mask
  }

  // ---------------------------------------------------------------- normalization

  /**
   * The official causal statistics: float64 running mean and Bessel standard deviation over the
   * observed steps (floored at 1e-6), rounded to float32. The cumulative sums run in order, as
   * NumPy's `cumsum` does (the first value starts the sum, so its signed zero is kept).
   */
  internal fun causalStatistics(
    x: FloatArray,
    observed: BooleanArray,
  ): Pair<FloatArray, FloatArray> {
    val n = x.size
    val means = FloatArray(n)
    val scales = FloatArray(n)
    var cumWeight = 0.0
    var cumValue = 0.0
    var m2 = 0.0
    var previousMean = 0.0
    for (t in 0 until n) {
      val value = x[t].toDouble()
      val weight = if (observed[t]) 1.0 else 0.0
      val weighted = weight * value
      if (t == 0) {
        cumWeight = weight
        cumValue = weighted
      } else {
        cumWeight += weight
        cumValue += weighted
      }
      val den = npMax(cumWeight, 1.0)
      val mean = cumValue / den
      val increment = (value - previousMean) * (value - mean) * weight
      m2 = if (t == 0) increment else m2 + increment
      val variance = m2 / npMax(den - 1.0, 1.0)
      means[t] = mean.toFloat()
      scales[t] = npMax(Math.sqrt(variance), MIN_SCALE).toFloat()
      previousMean = mean
    }
    return means to scales
  }

  /** Per step: the location and scale of its patch (the statistics at the patch's last step). */
  internal fun patchLocScale(
    x: FloatArray,
    observed: BooleanArray,
  ): Pair<FloatArray, FloatArray> {
    val (means, scales) = causalStatistics(x, observed)
    val loc = FloatArray(x.size)
    val scale = FloatArray(x.size)
    for (start in x.indices step PATCH) {
      val end = start + PATCH - 1
      loc.fill(means[end], start, start + PATCH)
      scale.fill(scales[end], start, start + PATCH)
    }
    return loc to scale
  }

  /** asinh((x - loc) / max(scale, eps)) in float32 on observed steps, 0 elsewhere. */
  internal fun normalize(
    x: FloatArray,
    loc: FloatArray,
    scale: FloatArray,
    observed: BooleanArray,
  ): FloatArray =
    FloatArray(x.size) { t ->
      if (observed[t]) {
        val y = (x[t] - loc[t]) / npMax(scale[t], F32_EPS)
        asinh(y.toDouble()).toFloat()
      } else {
        0f
      }
    }

  // ---------------------------------------------------------------- output

  /** clamp_nonfinite(sinh(preds) * scale + loc), all float32 but sinh. */
  internal fun denormalize(preds: FloatArray, loc: Float, scale: Float): FloatArray {
    val out = FloatArray(preds.size) { StrictMath.sinh(preds[it].toDouble()).toFloat() * scale + loc }
    clampNonFinite(out)
    return out
  }

  /**
   * In place, per step of 21 quantiles: +inf becomes the step's largest finite value, -inf its
   * smallest; an infinity left (no finite value in the step) becomes 0. NaN stays.
   */
  internal fun clampNonFinite(values: FloatArray) {
    for (start in values.indices step QUANTILES) {
      var hi = Float.NEGATIVE_INFINITY
      var lo = Float.POSITIVE_INFINITY
      for (i in start until start + QUANTILES) {
        val v = values[i]
        if (v.isFinite()) {
          hi = npMax(hi, v)
          lo = npMin(lo, v)
        }
      }
      for (i in start until start + QUANTILES) {
        var v = values[i]
        if (v == Float.POSITIVE_INFINITY) {
          v = hi
        } else if (v == Float.NEGATIVE_INFINITY) {
          v = lo
        }
        values[i] = if (v.isInfinite()) 0f else v
      }
    }
  }

  /**
   * The official `_refine_block` for one row of [h] decoder patches: the quantiles move from the
   * anchor frame (the last context patch's loc / scale) to the advancing frame of each patch.
   */
  internal fun refineBlock(
    quantiles: FloatArray,
    context: FloatArray,
    observed: BooleanArray,
    h: Int,
  ): FloatArray {
    if (h <= 1) return quantiles
    val (loc, scale) = patchLocScale(context, observed)
    val loc0 = loc[loc.size - 1]
    val scale0 = scale[scale.size - 1]
    val sinhMedian = DoubleArray(h * PATCH) { t ->
      (quantiles[t * QUANTILES + MEDIAN].toDouble() - loc0.toDouble()) / scale0.toDouble()
    }
    val frames = advancingBlockFrames(context, observed, sinhMedian, h)
    val out = FloatArray(quantiles.size)
    for (p in 0 until h) {
      val frameLoc = if (p == 0) loc0 else frames.first[p].toFloat()
      val frameScale = if (p == 0) scale0 else frames.second[p].toFloat()
      for (j in 0 until PATCH) {
        val base = (p * PATCH + j) * QUANTILES
        for (i in 0 until QUANTILES) {
          val sinhNorm = (quantiles[base + i] - loc0) / scale0
          out[base + i] = sinhNorm * frameScale + frameLoc
        }
      }
    }
    clampNonFinite(out)
    return out
  }

  /**
   * The official `advancing_block_frames` for one row: per decoder patch the float64 (loc, scale)
   * of a Welford state that starts from the whole context and takes in each decoded patch's median.
   */
  internal fun advancingBlockFrames(
    context: FloatArray,
    observed: BooleanArray,
    sinhMedian: DoubleArray,
    h: Int,
  ): Pair<DoubleArray, DoubleArray> {
    val contextPatches = context.size / PATCH
    val n = contextPatches + h
    val count = DoubleArray(n)
    val mean = DoubleArray(n)
    val m2 = DoubleArray(n)
    val products = DoubleArray(PATCH)
    val weights = DoubleArray(PATCH)
    for (p in 0 until contextPatches) {
      for (j in 0 until PATCH) {
        val t = p * PATCH + j
        weights[j] = if (observed[t]) 1.0 else 0.0
        products[j] = context[t].toDouble() * weights[j]
      }
      val c = reduceSum(weights)
      var m = reduceSum(products) / npMax(c, 1.0)
      for (j in 0 until PATCH) {
        val d = context[p * PATCH + j].toDouble() - m
        products[j] = weights[j] * (d * d)
      }
      var s = reduceSum(products)
      if (c == 0.0) {
        m = 0.0
        s = 0.0
      }
      count[p] = c
      mean[p] = m
      m2[p] = s
    }
    // The decoded patches start empty (count 0); only the scan's length depends on them.
    cumulativeWelford(count, mean, m2)
    var c = count[contextPatches - 1]
    var m = mean[contextPatches - 1]
    var s = m2[contextPatches - 1]
    val offset = m
    val locs = DoubleArray(h)
    val scales = DoubleArray(h)
    val folded = DoubleArray(PATCH)
    for (p in 0 until h) {
      val variance = npMax(s / npMax(c - 1.0, 1.0), 0.0)
      val sc = npMax(Math.sqrt(variance), MIN_SCALE)
      locs[p] = m
      scales[p] = sc
      for (j in 0 until PATCH) {
        val value = sinhMedian[p * PATCH + j] * sc + m - offset
        folded[j] = clip(value, -FOLD_CLIP, FOLD_CLIP) + offset
      }
      val blockMean = reduceSum(folded) / PATCH.toDouble()
      for (j in 0 until PATCH) {
        val d = folded[j] - blockMean
        products[j] = d * d
      }
      val blockM2 = reduceSum(products)
      val combined = welfordCombine(c, m, s, PATCH.toDouble(), blockMean, blockM2)
      c = combined[0]
      m = combined[1]
      s = combined[2]
    }
    return locs to scales
  }

  /** The official Welford merge of state a with state b: (w, mean, m2). */
  internal fun welfordCombine(
    wA: Double,
    meanA: Double,
    m2A: Double,
    wB: Double,
    meanB: Double,
    m2B: Double,
  ): DoubleArray {
    val w = wA + wB
    val wSafe = npMax(w, F64_TINY)
    val delta = meanB - meanA
    val mean = meanA + delta * wB / wSafe
    val m2 = m2A + m2B + delta * delta * wA * wB / wSafe
    return if (w == 0.0) doubleArrayOf(w, 0.0, 0.0) else doubleArrayOf(w, mean, m2)
  }

  /**
   * In place, the official Hillis-Steele scan of Welford states: after it, entry i holds the merge
   * of entries 0..i. Each round merges every entry with the one d before it (an empty state before
   * the start), d = 1, 2, 4 … while d < n.
   */
  internal fun cumulativeWelford(w: DoubleArray, mean: DoubleArray, m2: DoubleArray) {
    val n = w.size
    var d = 1
    while (d < n) {
      val w0 = w.copyOf()
      val mean0 = mean.copyOf()
      val m20 = m2.copyOf()
      for (i in 0 until n) {
        val wA = if (i >= d) w0[i - d] else 0.0
        val meanA = if (i >= d) mean0[i - d] else 0.0
        val m2A = if (i >= d) m20[i - d] else 0.0
        val merged = welfordCombine(wA, meanA, m2A, w0[i], mean0[i], m20[i])
        w[i] = merged[0]
        mean[i] = merged[1]
        m2[i] = merged[2]
      }
      d *= 2
    }
  }

  /**
   * Clamps a row's quantiles to [min - k s, max + k s] of its observed context, s the context's
   * last scale ([lastScale], float32).
   */
  internal fun cap(
    quantiles: FloatArray,
    rowValues: FloatArray,
    observed: BooleanArray,
    lastScale: Float,
  ): FloatArray {
    var contextMax = Float.NEGATIVE_INFINITY
    var contextMin = Float.POSITIVE_INFINITY
    for (t in rowValues.indices) {
      if (observed[t]) {
        contextMax = npMax(contextMax, rowValues[t])
        contextMin = npMin(contextMin, rowValues[t])
      }
    }
    if (contextMax == Float.NEGATIVE_INFINITY) contextMax = 0f
    if (contextMin == Float.POSITIVE_INFINITY) contextMin = 0f
    val anchor = lastScale * CAP_K
    val lo = contextMin - anchor
    val hi = contextMax + anchor
    return FloatArray(quantiles.size) { npMin(npMax(quantiles[it], lo), hi) }
  }

  private fun checkNoNan(values: FloatArray) {
    val nan = values.count { it.isNaN() }
    if (nan > 0) throw ForecastException("forecast produced $nan NaN values")
  }

  /** Whether every observed value is ≥ 0 (and at least one is observed). */
  internal fun nonNegativeContext(inputs: FloatArray, observed: BooleanArray): Boolean {
    var any = false
    for (t in inputs.indices) {
      if (observed[t] && inputs[t].isFinite()) {
        any = true
        if (!(inputs[t] >= 0f)) return false
      }
    }
    return any
  }

  // ---------------------------------------------------------------- NumPy arithmetic

  /**
   * NumPy's float64 sum of one axis of 8 to 128 values: eight running sums (values i, i + 8, …),
   * combined as ((r0 + r1) + (r2 + r3)) + ((r4 + r5) + (r6 + r7)), then the rest in order, added
   * to the reduction's start value 0.
   */
  internal fun reduceSum(values: DoubleArray): Double {
    val n = values.size
    require(n in 8..128) { "pairwise sum of $n values" }
    val r = DoubleArray(8) { values[it] }
    var i = 8
    while (i < n - n % 8) {
      for (j in 0 until 8) r[j] += values[i + j]
      i += 8
    }
    var sum = ((r[0] + r[1]) + (r[2] + r[3])) + ((r[4] + r[5]) + (r[6] + r[7]))
    while (i < n) {
      sum += values[i]
      i++
    }
    return 0.0 + sum
  }

  /**
   * np.maximum as NumPy 2.5 computes it, measured on every broadcast pattern and on reductions:
   * IEEE 754 maximum, NaN wins and -0.0 < +0.0 — the rule of [Math.max].
   */
  internal fun npMax(a: Double, b: Double): Double = Math.max(a, b)

  internal fun npMax(a: Float, b: Float): Float = Math.max(a, b)

  /** np.minimum: IEEE 754 minimum, NaN wins and -0.0 < +0.0 — the rule of [Math.min]. */
  internal fun npMin(a: Float, b: Float): Float = Math.min(a, b)

  /** np.clip: the value stays unless it is strictly outside [lo, hi] (so NaN and -0.0 stay). */
  internal fun clip(v: Float, lo: Float, hi: Float): Float =
    if (v < lo) lo else if (v > hi) hi else v

  internal fun clip(v: Double, lo: Double, hi: Double): Double =
    if (v < lo) lo else if (v > hi) hi else v

  /**
   * asinh in float64, fdlibm's `s_asinh.c` on [StrictMath] (Java has no asinh): accurate to about
   * one ulp, like the C library NumPy calls.
   */
  internal fun asinh(x: Double): Double {
    val bits = java.lang.Double.doubleToRawLongBits(x)
    val high = (bits ushr 32).toInt()
    val magnitude = high and 0x7fffffff
    if (magnitude >= 0x7ff00000) return x + x
    if (magnitude < 0x3e300000) return x
    val absX = Math.abs(x)
    val w =
      when {
        magnitude > 0x41b00000 -> StrictMath.log(absX) + LN2
        magnitude > 0x40000000 ->
          StrictMath.log(2.0 * absX + 1.0 / (Math.sqrt(x * x + 1.0) + absX))
        else -> {
          val t = x * x
          StrictMath.log1p(absX + t / (1.0 + Math.sqrt(1.0 + t)))
        }
      }
    return if (high > 0) w else -w
  }

  private const val LN2 = 6.93147180559945286227e-01
}