Skip to content

Use tfp-nightly for tests#2056

Closed
juanitorduz wants to merge 1 commit intopyro-ppl:masterfrom
juanitorduz:use-tfp-nightly
Closed

Use tfp-nightly for tests#2056
juanitorduz wants to merge 1 commit intopyro-ppl:masterfrom
juanitorduz:use-tfp-nightly

Conversation

@juanitorduz
Copy link
Copy Markdown
Collaborator

@juanitorduz juanitorduz commented Aug 2, 2025

Motivated by #2052 (comment) and in view of the comment tensorflow/probability#1994 (comment)

For TFP on JAX I would strongly advise using tfp-nightly. It's tested continuously against the latest JAX, and you can pin to a particular date's release if you need your dependencies to remain fixed.

(also by the fact that the last stable release of tensorflow-probability was like 9 months ago), I suggest using tfp-nightly in our test suite :)

It seems this one depends on #2052 (but maybe it's easier if this change is added there directly)

ImportError: cannot import name 'pjit_p' from 'jax.experimental.pjit'

@fehiepsi
Copy link
Copy Markdown
Member

fehiepsi commented Aug 4, 2025

Thank you @juanitorduz! @tare incorporated the change in #2052. The tests are passing now.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants