jax.debug package#

Runtime value debugging utilities#

jax.debug.print and jax.debug.breakpoint describes how to make use of JAX’s runtime value debugging features.

callback(callback, *args[, ordered])

Calls a stageable Python callback.

print(fmt, *args[, ordered])

Prints values and works in staged out JAX functions.

breakpoint(*[, backend, filter_frames, ...])

Enters a breakpoint at a point in a program.

Sharding debugging utilities#

Functions that enable inspecting and visualizing array shardings inside (and outside) staged functions.

inspect_array_sharding(value, *, callback)

Enables inspecting array sharding inside JIT-ted functions.

visualize_array_sharding(arr, **kwargs)

Visualizes an array's sharding.

visualize_sharding(shape, sharding, *[, ...])

Visualizes a Sharding using rich.