jaxlib.xla_extension.TpuDevice¶

class jaxlib.xla_extension.TpuDevice¶

Methods

`__init__`(args, *kwargs)
`live_buffers`(self)
`transfer_from_outfeed`(self, arg0)
`transfer_to_infeed`(self, arg0)

Attributes

`client`
`coords`	The coordinates of this TpuDevice's chip in the TPU mesh network.
`core_on_chip`	The index of this TpuDevice's core on the TPU chip.
`device_kind`
`host_id`	Deprecated; please use process_index
`id`	Integer ID of this device.
`platform`
`process_index`	Integer index of this device's process.
`task_id`	Deprecated; please use process_index