neilgirdhar / efax Goto Github PK

See wikipedia. Also, depends on #12

This would help with creating distributions like the normal-gamma distribution, and the normal inverse-Wishart distribution.

Depends on google/jax#2488

https://en.wikipedia.org/wiki/Inverse_Gaussian_distribution
https://docs.scipy.org/doc/scipy/reference/generated/scipy.stats.wald.html#scipy.stats.wald

Thanks for creating such a nice library.

I have experienced numerical instability in the Multivariate Gaussian Distributions. For better numerical stability, I wondered if it would be better replace e.g.,

With Choleksy based operations e.g.,

import jax.scipy as jsp

JITTER = 1e-6
half_precision = self.negative_half_precision
lower_h = jsp.linalg.cholesky(half_precision + JITTER  * jnp.eye(half_precision.shape[0]))
lower_h_inv = jsp.linalg.solve_triangular(lower_h,  jnp.eye(half_precision.shape[0]), lower=True)
h_inv = -jsp.linalg.solve_triangular(lower_h, lower_h_inv, lower=False)

We can also compute log determinant as ld = 2.0 * jnp.sum(jnp.log(jnp.diagonal(lower_h))). Downside though on using Cholesky decomposition in JAX, is it really wants you to be doing stuff in float64.

Depends on google/jax#2466

This is now well within reach thanks to th e Array API.

• For every method, find xp = get_namespace(*arrays) and then use xp instead of jax.numpy.
• Use the special extension for special functions (xp.special). Depends on data-apis/array-api#725.
• Support sampling methods for:
  • Jax
  • PyTorch
  • NumPy
• Support native fixed point sampling methods (used in exp-to-nat) for:
  • Jax
  • PyTorch
  • NumPy
• Generalize abstract_custom_jvp to PyTorch.
• Port the Fisher information code (which depends on automatic differentiation) to PyTorch.
• Move automatic JIT-application from methods to tests.
• Make tests work for each namespace
  • Jax
  • PyTorch
  • NumPy

https://en.wikipedia.org/wiki/Wishart_distribution
https://docs.scipy.org/doc/scipy/reference/generated/scipy.stats.wishart.html#scipy.stats.wishart

NEP 29 says to drop Python 3.8 on April 14, 2023. Use https://github.com/frostming/fix-future-annotations to upgrade the code.

Wait on python/mypy#14072