Suggest tfp-nighly

ricardoV94 · ricardoV94 · commit 40f5102d1d85 · 2025-12-01T14:46:19.000+01:00
diff --git a/pytensor/link/jax/dispatch/scalar.py b/pytensor/link/jax/dispatch/scalar.py
@@ -44,7 +44,7 @@ def try_import_tfp_jax_op(op: ScalarOp, jax_op_name: str | None = None) -> Calla
     except ModuleNotFoundError:
         raise NotImplementedError(
             f"No JAX implementation for Op {op.name}. "
-            "Implementation is available if TensorFlow Probability is installed"
+            "Implementation is available if tfp-nightly is installed"
         )
 
     if jax_op_name is None:

Original file line number	Diff line number	Diff line change
`@@ -44,7 +44,7 @@ def try_import_tfp_jax_op(op: ScalarOp, jax_op_name: str \| None = None) -> Calla`
`44`	`44`	`except ModuleNotFoundError:`
`45`	`45`	`raise NotImplementedError(`
`46`	`46`	`f"No JAX implementation for Op {op.name}. "`
`47`		`- "Implementation is available if TensorFlow Probability is installed"`
	`47`	`+ "Implementation is available if tfp-nightly is installed"`
`48`	`48`	`)`
`49`	`49`
`50`	`50`	`if jax_op_name is None:`