Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Import Jaxpr and ClosedJaxpr from jax.exend #544

Merged
merged 2 commits into from
Dec 19, 2024
Merged

Import Jaxpr and ClosedJaxpr from jax.exend #544

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
e794d9c
Import Jaxpr and ClosedJaxpr from jax.exend
jakevdp Dec 16, 2024
a5820a2
Fix pyright failures from numpy 2.
patrick-kidger Dec 17, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion diffrax/_misc.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -148,7 +148,7 @@ def static_select(pred: BoolScalarLike, a: ArrayLike, b: ArrayLike) -> ArrayLike
# This in turn allows us to perform some trace-time optimisations that XLA isn't
# smart enough to do on its own.
if isinstance(pred, (np.ndarray, np.generic)) and pred.shape == ():
pred = pred.item()
pred = cast(BoolScalarLike, pred.item())
if pred is True:
return a
elif pred is False:
Expand Down
5 changes: 3 additions & 2 deletions diffrax/_progress_meter.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -123,8 +123,9 @@ def _step_bar(bar: list[float], progress: FloatScalarLike) -> None:
if eqx.is_array(progress):
# May not be an array when called with `JAX_DISABLE_JIT=1`
progress = cast(Union[Array, np.ndarray], progress)
progress = progress.item()
progress = cast(float, progress)
progress = cast(float, progress.item())
else:
progress = cast(float, progress)
bar[0] = progress
print(f"{100 * progress:.2f}%")

Expand Down
4 changes: 2 additions & 2 deletions diffrax/_solver/runge_kutta.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@
import equinox as eqx
import equinox.internal as eqxi
import jax
import jax.core
import jax.extend as jex
import jax.lax as lax
import jax.numpy as jnp
import jax.tree_util as jtu
Expand Down Expand Up @@ -315,7 +315,7 @@ def _filter_stop_gradient(x):


def _is_jaxpr(x):
return isinstance(x, (jax.core.Jaxpr, jax.core.ClosedJaxpr))
return isinstance(x, (jex.core.Jaxpr, jex.core.ClosedJaxpr))


def _filter_maybe_cond(pred, branch, value):
Expand Down
Loading