betaincinv skips the root finding when y == 0.5, so Beta medians are wrong

未关闭
  • 哪里不对？
#2,034 0 条评论 0 个 reaction 已指派 0 人 在 GitHub 查看

还没有人认领这个 Issue。

评估

难度
3/5
预计耗时
1-2 天
新手友好度
68/100
Issue 类型
缺陷
描述清晰度
描述清楚
活跃度
活跃
技术栈
python
领域
backend

调研方向

Start in tensorflow_probability/python/math/special.py at _betaincinv_computation and run the issue's reproduction to observe the incorrect median results. Add a regression test for y == 0.5, including the symmetric Beta case, and verify that the quantile results are correct without regressing other inputs.

由索引模型根据 Issue 内容生成。

描述

tfp.math.betaincinv(a, b, y) returns its initial approximation, without any root-finding iterations, when y == 0.5. So every median computed through it is off, including tfd.Beta(a, b).quantile(0.5). In float64 this happens for all a, b. In float32 it happens for part of the a < 1, b < 1 region, where the bracket is halved: (0.5, 0.5) gives 0.6185 and (0.9, 0.5) gives 0.7783, against 0.5 and 0.7210.

System info
  • tfp-nightly 0.26.0.dev20261008 with jax 0.10.2, and tensorflow-probability 0.25.0 with jax 0.4.30. Same output on both.
  • Linux (WSL2), CPU. The script downloads nothing.
Reproduction
import jax

jax.config.update("jax_enable_x64", True)
import jax.numpy as jnp
import scipy.stats as st
from tensorflow_probability.substrates import jax as tfp

for a, b in [(2.0, 3.0), (10.0, 1.0), (0.5, 0.5)]:
    got = tfp.math.betaincinv(jnp.float64(a), jnp.float64(b), jnp.float64(0.5))
    quantile = tfp.distributions.Beta(jnp.float64(a), jnp.float64(b)).quantile(jnp.float64(0.5))
    near = tfp.math.betaincinv(jnp.float64(a), jnp.float64(b), jnp.float64(0.5 - 1e-12))
    print(f"a={a} b={b}  betaincinv(y=0.5)={float(got):.6f}  Beta.quantile(0.5)={float(quantile):.6f}  "
          f"betaincinv(y=0.5-1e-12)={float(near):.6f}  scipy={st.beta(a, b).ppf(0.5):.6f}")

Current result:

a=2.0 b=3.0  betaincinv(y=0.5)=0.390087  Beta.quantile(0.5)=0.390087  betaincinv(y=0.5-1e-12)=0.385728  scipy=0.385728
a=10.0 b=1.0  betaincinv(y=0.5)=0.906306  Beta.quantile(0.5)=0.906306  betaincinv(y=0.5-1e-12)=0.933033  scipy=0.933033
a=0.5 b=0.5  betaincinv(y=0.5)=0.381514  Beta.quantile(0.5)=0.381514  betaincinv(y=0.5-1e-12)=0.499537  scipy=0.500000

(Beta(10, 1) has median 0.5 ** (1 / 10) = 0.933033.)

Root cause

In _betaincinv_computation the while loop starts with this stop mask (special.py#L753):

      loop_vars=(
          tf.equal(y, initial_low) | tf.equal(y, initial_high),
          initial_low,
          initial_high,
          initial_candidate),

initial_low and initial_high bracket the solution x, not the probability y. In float64 initial_high is always 0.5 (in float32 it is 0.5 in the small a, b branch), so y == 0.5 marks the element as finished before the first iteration and initial_candidate comes back as the result. The y == 0 and y == 1 cases are already handled after the loop, so I think the intended mask is:

          tf.equal(y, zero) | tf.equal(y, one),

With that one-line change the first two rows give 0.385728 and 0.933033, matching scipy. The mask only differs when y is 0.5, so other outputs are unchanged: 4000 random (a, b, y) triples came out bitwise identical before and after.

A second, smaller problem

For a == b the root at y = 0.5 sits right at the 0.5 bracket edge. Halley steps land just past high, the bisection fallback only halves the remaining gap, and after the 8 float64 iterations (0.5, 0.5) ends at 0.499537. This is the y=0.5-1e-12 column above, which goes through the full loop today. The one-line fix doesn't change this. Returning 0.5 directly when betainc(a, b, 0.5) is within tolerance of y would cover it, but I haven't tested that.

Downstream

NumPyro's Beta.icdf and StudentT.icdf call this function. StudentT.icdf passes 1 - |1 - 2q|, which is 0.5 at both quartiles, so in x64 dist.StudentT(3.0).icdf(0.75) gives 0.664628 instead of 0.764892. TFP's own StudentT.quantile isn't affected, since it goes through stdtrit.

I searched the issues for betaincinv, Beta quantile, StudentT quantile and quantile 0.5. The closest are #1581 and #1625 (2022, which added the current implementation), and neither mentions this. Happy to open a PR with the change and a regression test if that would help.

I used Claude Code to help trace this. I ran the reproduction and checked the numbers myself.

主要语言
Jupyter Notebook
星标
4.4k
派生
1.1k
30 天内没有已合并 PR

环境准备

  • 没有 Dockerfile 或 Docker Compose 文件
  • 没有 Pull Request 模板
  • 阅读贡献指南

从这里开始

  1. 先读完整个 Issue，再读项目的贡献指南。
  2. 在 Issue 下留言说明你要接手 —— 这能避免两个人做同样的事。
  3. Fork 仓库，在一个分支上完成修改。
  4. 提交 Pull Request，并在描述里引用这个 Issue 编号。

tensorflow/probability 的其他 Issue

查看 tensorflow/probability 的全部 Issue

相似的 Issue

更多 Backend & API Design Issue

把新 issue 发到你的邮箱

精选适合新手参与的 GitHub issue 摘要。