jax.experimental.pallas.mosaic_gpu.query_cluster_cancel#

jax.experimental.pallas.mosaic_gpu.query_cluster_cancel(result_ref, grid_names)[源代码]#

解码 try_cluster_cancel 操作的结果。

它解释由已完成的 try_cluster_cancel 调用写入共享内存的 16 字节不透明响应，以确定是否成功获取了新的工作单元。

参数:

  • result_ref (_Ref) – 包含查询响应的 SMEM ref。

  • grid_names (Sequence[Hashable]) – 要查询的网格轴名称元组。

返回:

  • 所请求轴名称的网格索引。

  • 一个布尔值，指示取消是否成功。

返回类型:

包含解码响应的元组

另请参阅

jax.experimental.pallas.mosaic_gpu.try_cluster_cancel()