jax.experimental.pallas.mosaic_gpu.SemaphoreSignal#

class jax.experimental.pallas.mosaic_gpu.SemaphoreSignal(ref: '_Ref', *, device_id: 'pallas_primitives.DeviceId | None', inc: 'int | jax.Array' = 1)[source]#
Parameters:

  • ref (_Ref)

  • device_id (pallas_primitives.DeviceId | None)

  • inc (int | jax.Array)

__init__(ref, *, device_id, inc=1)#
Parameters:

  • ref (_Ref)

  • device_id (pallas_primitives.DeviceId | None)

  • inc (int | jax.Array)

Return type:

None

Methods

__init__(ref, *, device_id[, inc])

Attributes

inc

ref

device_id