jax.experimental.pallas.tpu.SemaphoreType#

class jax.experimental.pallas.tpu.SemaphoreType(value, names=<not given>, *values, module=None, qualname=None, type=None, start=1, boundary=None)[source]#
__init__(*args, **kwds)#

Methods

get_array_aval()

get_ref_aval()

Attributes

REGULAR

DMA

BARRIER