jax.debug.inspect_array_sharding#
- jax.debug.inspect_array_sharding(value, *, callback)[source]#
启用在 JIT 编译函数内部检查数组分片。
此函数在提供数组的 Pytree 时,会回调每个数组的分片并在
pjit
编译的计算中工作,从而能够检查所选的中间分片。当分片信息可用时,
callback
被调用的策略是尽早。这意味着如果在没有任何转换的情况下调用inspect_array_callback
,则回调会立即发生,因为我们已经拥有数组及其分片。在jax.jit
内部,回调将在降低时发生,这意味着您可以使用AOT API(jit(f).lower(...)
)触发回调。在pjit
内部,回调在编译时发生,因为分片由XLA确定。您可以通过使用JAX的AOT API(pjit(f).lower(...).compile()
)触发回调。在所有情况下,回调都将通过运行函数来触发,因为运行函数首先需要降低和编译它。但是,一旦函数被编译并缓存,回调将不再发生。此函数为实验性功能,其行为将来可能会发生变化。
- 参数:
value – JAX数组的Pytree。
callback (Callable[[Sharding], None]) – 一个可调用对象,它接收一个
Sharding
并且不返回值。
在以下示例中,我们打印出
pjit
计算中中间值的碎片。>>> import jax >>> import jax.numpy as jnp >>> from jax.experimental.pjit import pjit >>> from jax.sharding import Mesh, PartitionSpec >>> >>> x = jnp.arange(8, dtype=jnp.float32) >>> def f_(x): ... x = jnp.sin(x) ... jax.debug.inspect_array_sharding(x, callback=print) ... return jnp.square(x) >>> f = pjit(f_, in_shardings=PartitionSpec('dev'), ... out_shardings=PartitionSpec('dev')) >>> with Mesh(jax.devices(), ('dev',)): ... f.lower(x).compile() ... NamedSharding(mesh={'dev': 8}, partition_spec=PartitionSpec(('dev',),))