Relax tensorflow-probability dependency

[sritchie]

This PR:

• fixes 3.13 installation issues by pinning to sub-2.0 poetry (EDIT: moved to Downgrade poetry (GEN-990) #1549 )
• relaxes our tensorflow-probability dependency, so that colab installation doesn't require a downgrade of tfp
• updates the dependencies in our lockfile

This upgrade exposed two bugs... one's linked in the comment below, another's filed here: jax-ml/jax#26908

[linear]

GEN-990 [BUG] Cannot install dependencies under Python 3.13 on MacOS (ARM)

Describe the bug
Installing dependencies for this repo using poetry install fails for Python 3.13. (But works for Python 3.12.)

To Reproduce

• Clone this repo
• Run poetry install as described in https://github.com/ChiSym/genjax/blob/main/CONTRIBUTING.md
• See an error when installing scipy:

    | FAILED: scipy/_lib/_uarray/_uarray.cpython-313-darwin.so.p/_uarray_dispatch.cxx.o 
    | /opt/homebrew/opt/llvm/bin/clang++ -Iscipy/_lib/_uarray/_uarray.cpython-313-darwin.so.p -Iscipy/_lib/_uarray -I../scipy/_lib/_uarray -I/opt/homebrew/opt/python3.13/Frameworks/Python.framework/Versions/3.13/include/python3.13 -I/opt/homebrew/opt/llvm/include -fvisibility=hidden -fvisibility-inlines-hidden -fdiagnostics-color=always -DNDEBUG -Wall -Winvalid-pch -std=c++17 -O3 -Wno-unused-function -MD -MQ scipy/_lib/_uarray/_uarray.cpython-313-darwin.so.p/_uarray_dispatch.cxx.o -MF scipy/_lib/_uarray/_uarray.cpython-313-darwin.so.p/_uarray_dispatch.cxx.o.d -o scipy/_lib/_uarray/_uarray.cpython-313-darwin.so.p/_uarray_dispatch.cxx.o -c ../scipy/_lib/_uarray/_uarray_dispatch.cxx
    | In file included from ../scipy/_lib/_uarray/_uarray_dispatch.cxx:3:
    | ../scipy/_lib/_uarray/small_dynamic_array.h:145:18: error: reference to non-static member function must be called
    |   145 |     size_ = copy.size;
    |       |             ~~~~~^~~~
    | 1 error generated.
    | [52/1383] Compiling C object scipy/special/libcdflib.a.p/cdflib.c.o
    | [53/1383] Compiling C++ object scipy/special/_ufuncs_cxx.cpython-313-darwin.so.p/_faddeeva.cxx.o
    | [54/1383] Compiling C++ object scipy/special/_ufuncs.cpython-313-darwin.so.p/sf_error.cc.o
    | [55/1383] Compiling C++ object scipy/special/_ufuncs_cxx.cpython-313-darwin.so.p/_wright.cxx.o
    | [56/1383] Compiling C++ object scipy/special/_special_ufuncs.cpython-313-darwin.so.p/sf_error.cc.o
    | [57/1383] Compiling C++ object scipy/special/_gufuncs.cpython-313-darwin.so.p/sf_error.cc.o
    | [58/1383] Compiling C object scipy/_lib/_test_deprecation_def.cpython-313-darwin.so.p/meson-generated__test_deprecation_def.c.o
    | [59/1383] Generating 'scipy/_lib/_ccallback_c.cpython-313-darwin.so.p/_ccallback_c.c'
    | [60/1383] Compiling C++ object scipy/special/_ufuncs.cpython-313-darwin.so.p/dd_real_wrappers.cpp.o
    | [61/1383] Compiling C object scipy/_lib/_test_deprecation_call.cpython-313-darwin.so.p/meson-generated__test_deprecation_call.c.o
    | [62/1383] Generating 'scipy/_lib/messagestream.cpython-313-darwin.so.p/messagestream.c'
    | [63/1383] Generating 'scipy/special/_specfun.cpython-313-darwin.so.p/_specfun.cpp'
    | [64/1383] Generating 'scipy/special/_ufuncs_cxx.cpython-313-darwin.so.p/_ufuncs_cxx.cpp'
    | [65/1383] Compiling C++ object scipy/special/_gufuncs.cpython-313-darwin.so.p/_gufuncs.cpp.o
    | [66/1383] Compiling C++ object scipy/special/_ufuncs.cpython-313-darwin.so.p/special_wrappers.cpp.o
    | [67/1383] Generating 'scipy/special/_ufuncs.cpython-313-darwin.so.p/_ufuncs.c'
    | [68/1383] Compiling C++ object scipy/special/_special_ufuncs.cpython-313-darwin.so.p/_special_ufuncs.cpp.o
    | ninja: build stopped: subcommand failed.

I played around with installing OpenBLAS manually with brew and updating LLVM, but I couldn't get it to work.

Expected behavior
Dependencies are installed cleanly.

Execution environment (please complete the following information):

• OS: macOS Sequoia 15.3.1
• Python version: 3.13.1
• GenJAX version: latest (54df069)

Workaround
Installing dependencies worked after I ran poetry env use python3.12

Possible fix
Either document the supported Python versions, or even better, specify the supported versions in the pyproject.toml.

[sritchie]

@femtomc , new_main (which we use in forward.py https://github.com/ChiSym/genjax/blob/main/src/genjax/_src/core/interpreters/forward.py#L41-L47) disappeared from JAX here:

jax-ml/jax@c36e1f7#diff-713d667c14b1dd075ec025dbd1cd8ea222e42c8b4c094ae05373d4d05b276dc5

can you help me figure out how to fix it?

[gitstream-cm]

🥷 Code experts: femtomc

sritchie, MathieuHuot have most 👩‍💻 activity in the files.
sritchie, femtomc have most 🧠 knowledge in the files.

See details

.pre-commit-config.yaml

Activity based on git-commit:

	sritchie	MathieuHuot
MAR
FEB	0 additions & 1 deletions
JAN	1 additions & 1 deletions
DEC	2 additions & 0 deletions
NOV
OCT

Knowledge based on git-blame:
sritchie: 85%

poetry.lock

Activity based on git-commit:

	sritchie	MathieuHuot
MAR
FEB	60 additions & 16 deletions
JAN	19 additions & 81 deletions
DEC
NOV
OCT	4 additions & 4 deletions

Knowledge based on git-blame:
femtomc: 32%
sritchie: 8%

pyproject.toml

Activity based on git-commit:

	sritchie	MathieuHuot
MAR	1 additions & 1 deletions
FEB	11 additions & 10 deletions
JAN	6 additions & 11 deletions
DEC
NOV
OCT	1 additions & 1 deletions

Knowledge based on git-blame:
sritchie: 30%
femtomc: 13%

tests/generative_functions/test_vmap_combinator.py

Activity based on git-commit:

	sritchie	MathieuHuot
MAR
FEB	20 additions & 1 deletions	57 additions & 1 deletions
JAN	5 additions & 1 deletions
DEC
NOV	29 additions & 0 deletions
OCT	59 additions & 17 deletions

Knowledge based on git-blame:
sritchie: 51%
femtomc: 8%

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

[femtomc]

We can actually expunge a lot of that code, and I’ve done so on my delayed sampling branch if you want to copy.