google-research / google-research/vmoe
Replace jax.scipy.stats.norm.cdf with jax.scipy.special.ndtr for improved efficiency
- Dominant language
- Jupyter Notebook
- Stars
- 728
- Forks
- 58
- Avg merge
- 4h 9m
- Merged PRs (30d)
- 3
Description
https://github.com/google-research/vmoe/blob/1e207e8a0fc8640bc5aab1fdd4b4fe6957d6c2a1/vmoe/nn/routing.py#L179
Current Code:
`p = 1. - jax.scipy.stats.norm.cdf(noise_required_to_win)`
Recommended Replacement:
```
from jax.scipy.special import ndtr
p = 1. - ndtr(noise_required_to_win)
```
The current code uses jax.scipy.stats.norm.cdf, which is based on a frozen distribution object (rv_continuous). This approach is flexible but adds unnecessary overhead for computing the standard normal cumulative distribution function (CDF), especially in performance-critical JAX pipelines.
In contrast, jax.scipy.special.ndtr is a low-level, efficient implementation of the standard normal CDF. It avoids object instantiation and works better with JAX’s JIT compilation and vectorized operations.
Both implementations are numerically equivalent, but ndtr is faster and more appropriate when only the standard normal distribution is needed.
Contributor guide
Assessment
This issue has not been assessed yet.