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