jax.experimental.pallas.GridSpec#

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

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

有关参数的更详细说明，请参阅jax.experimental.pallas.pallas_call()的文档，以及网格和BlockSpecs。

参数:

  • 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