google-research / google-research/vmoe

Replace jax.scipy.stats.norm.cdf with jax.scipy.special.ndtr for improved efficiency

Open
#200 0 comments 0 reactions 0 assignees View on GitHub
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.