ChipVersion(value[, names, module, ...])

CompilerParams([dimension_semantics, ...])

Mosaic TPU compiler parameters.

GridDimensionSemantics(value[, names, ...])

MemorySpace(value[, names, module, ...])

PrefetchScalarGridSpec(num_scalar_prefetch)

SemaphoreType(value[, names, module, ...])

TpuInfo(*, chip_version, generation, ...[, ...])

TPU hardware information.

load(ref, *[, mask])

Loads an array from the given ref.

store(ref, val, *[, mask])

Stores a value to the given ref.

async_copy(src_ref, dst_ref, sem, *[, ...])

Issues a DMA copying from src_ref to dst_ref.

async_remote_copy(src_ref, dst_ref, ...[, ...])

Issues a remote DMA copying from src_ref to dst_ref.

make_async_copy(src_ref, dst_ref, sem)

Creates a description of an asynchronous copy operation.

make_async_remote_copy(src_ref, dst_ref, ...)

Creates a description of a remote copy operation.

sync_copy(src_ref, dst_ref, *[, add])

Synchronously copies a PyTree of refs to another PyTree of refs.

BufferedRef(_spec, _buffer_type, window_ref, ...)

A helper class to automate VMEM double buffering in pallas pipelines.

BufferedRefBase()

Abstract interface for BufferedRefs.

emit_pipeline(body, *, grid[, in_specs, ...])

Creates a function to emit a manual pallas pipeline.

emit_pipeline_with_allocations(body, *, grid)

Creates pallas pipeline and top-level allocation preparation functions.

get_pipeline_schedule(schedule)

Retrieve a named pipeline schedule or pass through fully specified one.

make_pipeline_allocations(*refs[, in_specs, ...])

Create BufferedRefs for the pipeline.

prng_seed(*seeds)

Sets the seed for PRNG.

sample_block(sampler_fn, global_key, ...[, ...])

Samples a block of random values with invariance guarantees.

stateful_bernoulli(*args, **kwargs)

Sample Bernoulli random values with given shape and mean.

stateful_bits(*args, **kwargs)

Sample uniform bits in the form of unsigned integers.

stateful_normal(*args, **kwargs)

Sample standard normal random values with given shape and float dtype.

stateful_uniform(*args, **kwargs)

Sample uniform random values in [minval, maxval) with given shape/dtype.

to_pallas_key(key)

Helper function for converting non-Pallas PRNG keys into Pallas keys.

force_tpu_interpret_mode([params])

Context manager that forces TPU interpret mode under its dynamic context.

InterpretParams(*[, detect_races, ...])

Parameters for TPU interpret mode.

reset_tpu_interpret_mode_state()

Resets all global, shared state used by TPU interpret mode.

set_tpu_interpret_mode([params])

core_barrier(sem, *, core_axis_name)

Synchronizes all cores in a given axis.

get_barrier_semaphore()

Returns a barrier semaphore.

get_tpu_info()

Returns the TPU hardware information for the current device.

is_tpu_device()

Returns whether the current device is a TPU.

run_on_first_core(core_axis_name)

Runs a function on the first core in a given axis.

with_memory_space_constraint(x, memory_space)

Constrains the memory space of an array.