jax.debug 模块#

运行时值调试实用程序#

编译后的打印和断点 描述了如何使用 JAX 的运行时值调试功能。

callback(callback, *args[, ordered])

调用可分阶段的 Python 回调。

print(fmt, *args[, ordered])

打印值,并在分阶段的 JAX 函数中工作。

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

在程序中的某个点进入断点。

分片调试实用程序#

使您能够检查和可视化分阶段函数内部(和外部)的数组分片的功能。

inspect_array_sharding(value, *, callback)

在 JIT 编译函数中启用对数组分片进行检查。

visualize_array_sharding(arr, **kwargs)

可视化数组的分片。

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

使用 rich 可视化 Sharding