jax.debug.inspect_array_sharding

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',),))