Limarta/beta bug

[limarta]

This PR corrects the GEM's logpdf by solving for the correct beta values.

Fixes https://linear.app/chi-fro/issue/EXPERT-62/formula-for-beta-values-incorrect

[linear]

EXPERT-62 Formula for `beta` values incorrect?

In logpdf here, passing weights can sometimes produce negative beta values which is outside of the range [0,1]. This line seems responsible:

 betas = jax.vmap(lambda i: 1 - pi[i] / pi[i-1])(jnp.arange(len(pi)))

It isn't clear to me that when given a pi, the beta values can be solved for in parallel (e.g. vmap), so instead we use a lax.scan:

def recover(carry, x):
  beta = x / carry
  return carry * (1-beta), beta

jax.lax.scan(recover, 1.0, weights)

Timeline: Get it merged end of week.

[limarta]

@Joaoloula When you have a chance, could you confirm that the beta values are calculated here correctly? I think this may need to be fixed in the pure jax branch as well.

[gitstream-cm]

🥷 Code experts: Joaoloula, Schaechtle

Schaechtle has most 👩‍💻 activity in the files.
Joaoloula has most 🧠 knowledge in the files.

See details

pyproject.toml

Activity based on git-commit:

	Schaechtle
FEB
JAN
DEC
NOV
OCT	1 additions & 0 deletions
SEP

Knowledge based on git-blame:
Joaoloula: 93%

requirements-dev.lock

Activity based on git-commit:

	Schaechtle
FEB
JAN
DEC
NOV
OCT	2 additions & 0 deletions
SEP

Knowledge based on git-blame:
Joaoloula: 63%

requirements.lock

Activity based on git-commit:

	Schaechtle
FEB
JAN
DEC
NOV
OCT	2 additions & 0 deletions
SEP

Knowledge based on git-blame:
Joaoloula: 63%

src/genspn/distributions.py

Activity based on git-commit:

	Schaechtle
FEB
JAN
DEC
NOV
OCT	62 additions & 51 deletions
SEP

Knowledge based on git-blame:
Joaoloula: 66%

test/test_score.py

Activity based on git-commit:

	Schaechtle
FEB
JAN
DEC
NOV
OCT	28 additions & 33 deletions
SEP

Knowledge based on git-blame:
Joaoloula: 79%

To learn more about /:\ gitStream - Visit our Docs