Drop tensorflow_probability so scenvi imports on current jax

[Marius1311]

Motivation

pip install scenvi currently gives you a package that cannot be imported at all.
tensorflow_probability is pinned to ^0.22.0 — i.e. <0.23 — and tfp 0.22 fails
against any recent jax:

AttributeError: module 'jax.interpreters.xla' has no attribute 'pytype_aval_mappings'

This is what is failing every build job on this repo right now, including on
#91 and #92, where collection errors out before any test runs.

Re-pinning buys little: tfp's last release is 0.25.0 (November 2024) and it is
unmaintained, while jax is at 0.11. The dependency is also much larger than what
scenvi uses — four log-densities and one matrix helper, all closed-form — so this
removes it rather than chasing a compatible window.

What this does

• scenvi/_dists.py implements Poisson, negative binomial, zero-inflated negative
  binomial and unit-variance normal log-densities against
  jax.scipy.special.gammaln / jax.nn.log_sigmoid, plus fill_triangular.
• scenvi/ENVI.py loses its tensorflow_probability import and calls the local
  fill_triangular.
• tests/test_dists.py checks each replacement against the tfp original to
  rtol=1e-11 in float64, including saturated logits and large counts. It skips
  where tfp is absent — after this PR, everywhere except a deliberately pinned
  environment — and exists as the record that the replacements reproduce what they
  replaced.

Two details that the differential test caught, and that are easy to get wrong:

• tfp builds Inflated's mixture from categorical logits [d, -d], whose
  difference is 2d. The zero-inflation weight is therefore sigmoid(2d), not
  sigmoid(d) — using the latter silently rescales the parameter.
• fill_triangular's fill order is not the obvious one: the last m - n entries
  are laid down first, then the whole vector reversed on top, then the lower
  triangle taken. [1..6] becomes [[4,0,0],[6,5,0],[3,2,1]].

The flax pin has gone the same way

Removing tfp is necessary but not sufficient. flax = "^0.10.4" caps below 0.11,
and flax 0.10.7 calls jax.core.get_opaque_trace_state, which jax 0.11 removed —
so a fresh install still fails, just later and for a different reason. flax, optax
and clu therefore lose their upper caps, and jax becomes a direct dependency
instead of one inherited from flax.

The new bounds are written >= rather than ^ deliberately: poetry's caret on a
0.x version caps at the next minor, which is precisely how these pins became
stale. ^0.22.0 meant <0.23; ^0.10.4 means <0.11.

Testing

environment	result
clean install → jax 0.11.0, flax 0.12.8, optax 0.2.8, clu 0.0.12, no tfp	import scenvi and scenvi.ENVI work; 8 passed, 1 skipped
pinned jax 0.4.23 + tfp 0.22, where tfp still imports	19 passed, including the 11 differential tests

The first row is the case that is broken on main today.

[Marius1311]

CI is green here across 3.10–3.13. Worth noting for merge order: #91 and #92 will keep showing red build jobs until this one lands, since they don't touch the pins — their failures are this same collection error, not anything in those PRs.