jax.experimental.pallas.GridSpec#

类 jax.experimental.pallas.GridSpec(grid=(), in_specs=NoBlockSpec, out_specs=NoBlockSpec, scratch_shapes=())[源]#

编码 jax.experimental.pallas.pallas_call() 的网格参数。

请参阅 jax.experimental.pallas.pallas_call() 的文档，以及 网格和块规格 以获取更详细的参数说明。

参数:

  • grid (TupleGrid)

  • in_specs (BlockSpecTree)

  • out_specs (BlockSpecTree)

  • scratch_shapes (ScratchShapeTree)

__init__(grid=(), in_specs=NoBlockSpec, out_specs=NoBlockSpec, scratch_shapes=())[源]#
参数:

  • grid (Grid)

  • in_specs (BlockSpecTree)

  • out_specs (BlockSpecTree)

  • scratch_shapes (ScratchShapeTree)

方法

__init__([grid, in_specs, out_specs, ...])

属性

scratch_shapes

grid

grid_names

in_specs

out_specs