callback(callback, *args[, ordered, partitioned])

Calls a stageable Python callback.

print(fmt, *args[, ordered, partitioned, ...])

Prints values and works in staged out JAX functions.

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

Enters a breakpoint at a point in a program.

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 aSharding usingrich.