API

Look up the functions you call from Python here. The signatures and detailed descriptions below are generated from the source. For a complete program, use the User's Guide; for a reminder of which operations belong on the host or in a kernel, use the quick reference.

Task	Call	Result or effect
Choose the default device	tack.set_default_device("cuda")	Arrays are allocated there; with nothing set it is the CPU
Name a device	tack.device("cuda"), with tack.device("cuda"):	The same Device every time; with makes it the default in the block
Allocate on a device	tack.empty((n,), device="cuda")	Every constructor takes device=
Move an array	a.to_device("cpu"), a.device	a itself if its memory is usable there, else a copy
Check indices	tack.set_debug(True)	Every kernel launched afterwards checks its indices
Allocate an output	tack.empty((n,))	New f32 array with unspecified contents
Allocate with another dtype	tack.empty((n,), tack.i32)	Any Tack type or NumPy dtype (np.int32, "int32")
Allocate zeros, ones or a value	tack.zeros((n,)), tack.ones((n,)), tack.full((n,), 2.5)	New array, filled on the device
Match storage without copying	tack.empty_like(prototype)	Same shape and element dtype; no values read or initialized
Match and fill	tack.zeros_like(a), tack.ones_like(a), tack.full_like(a, v)	Same shape and element dtype, filled
Upload NumPy data	tack.array(host)	New array with an independent copy; a NumPy array keeps its dtype
Use an array as is	tack.asarray(obj)	obj itself if it is already a Tack array (of dtype, if given), else tack.array(obj)
Overwrite values from the host	tack.copyto(a, host)	Copies into a; host has a.numpy()'s shape or broadcasts to it
Copy between arrays	tack.copyto(dst, src)	Device-side copy into an existing array of the same shape
Download data	a.numpy() or np.asarray(a)	Independent NumPy array; waits for earlier writes
Allocate vector elements	tack.empty((n,), tack.vec3())	Array of n 3-vectors; NumPy download shape (n, 3)
Inspect shape and size	a.shape, a.ndim, a.size, a.nbytes	Element counts (components excluded) and storage bytes
Reshape without copying	a.reshape(rows, -1)	Shares storage; one dimension may be -1
Wait for device work	tack.sync()	Waits for queued launches on every device; returns None
Inspect a kernel	tack.inspect(kernel, *args, mode="source")	Generated source as a string; does not execute the kernel

Without a dtype, an allocation is f32; values from Python make f32 from floats and i32 from ints (i64 if a value needs it), and a NumPy array keeps its own dtype, so a float64 array becomes f64.

For ownership, errors and precise guarantees, see the Runtime API contract and the Kernel Language Contract. Compiler internals are in the Developer's Guide.

Devices

A kernel runs on its arrays' device: the default device when that can address them, otherwise theirs. See Backends: Devices.

tack.device

device(spec='cpu', **options) -> Device

The device spec names: "cpu", "metal", "cuda", "cuda:0", "hip", "level_zero"; a Device is returned as it is.

Without options it is made on first use and the same device afterwards, so its compiled kernels and arrays are shared by everything that asks. Options go to the backend and must be ones it declares in Backend.device_options; they make a new device each call. The CPU takes num_threads; Level Zero takes external_context, the driver, device and context handles of a context owned by another library, so both can address the same memory.

"cuda:N" is device N as the runtime numbers it (CUDA's ordinal, HIP's device id, the index among the first Level Zero driver's devices); device_count says how many there are. A bare "cuda" or "hip" is the device the calling thread already uses: the one of a CUDA context another library made current, HIP's current device, else device 0. "level_zero" is device 0.

with tack.device("cuda"): makes it the default device in the block.

Source code in packages/tack-core/src/tack/runtime/dispatch.py

def device(spec="cpu", **options) -> Device:
    """The device ``spec`` names: ``"cpu"``, ``"metal"``, ``"cuda"``, ``"cuda:0"``,
    ``"hip"``, ``"level_zero"``; a `Device` is returned as it is.

    Without options it is made on first use and the same device afterwards,
    so its compiled kernels and arrays are shared by everything that asks.
    Options go to the backend and must be ones it declares in
    `Backend.device_options`; they make a new device each call. The CPU
    takes ``num_threads``; Level Zero takes ``external_context``, the
    driver, device and context handles of a context owned by another
    library, so both can address the same memory.

    ``"cuda:N"`` is device N as the runtime numbers it (CUDA's ordinal,
    HIP's device id, the index among the first Level Zero driver's
    devices); `device_count` says how many there are. A bare ``"cuda"`` or
    ``"hip"`` is the device the calling thread already uses: the one of a
    CUDA context another library made current, HIP's current device, else
    device 0. ``"level_zero"`` is device 0.

    ``with tack.device("cuda"):`` makes it the default device in the block.
    """
    if isinstance(spec, Device):
        if options:
            raise TypeError("options make a new device: pass a device spec, not a Device")
        return spec
    kind, index = _parse(spec)
    if options:
        return as_device(_construct(kind, options, index))
    if index is None:
        index = _current_index(kind)
    with _lock:
        found = _devices.get((kind, index))
        if found is None:
            found = _devices[(kind, index)] = _construct(kind, {}, index)
        return as_device(found)

tack.Device

Device(backend)

A device: where arrays live and kernels run. tack.device(...) makes one.

type is "cpu", "metal", "cuda", "hip" or "level_zero"; index numbers the devices of a type. Two devices are equal when they are the same backend. with device: makes it the default device inside the block, on this thread.

Source code in packages/tack-core/src/tack/runtime/dispatch.py

def __init__(self, backend):
    self._backend = backend

synchronize

synchronize()

Wait for every kernel launched on this device so far.

Source code in packages/tack-core/src/tack/runtime/dispatch.py

def synchronize(self):
    """Wait for every kernel launched on this device so far."""
    self._backend.synchronize()

tack.device_count

device_count(kind: str) -> int

How many devices of type kind this machine has: 0 where it cannot start.

Counting makes no device: tack.device(f"cuda:{n}") for n below the count does.

Source code in packages/tack-core/src/tack/runtime/dispatch.py

def device_count(kind: str) -> int:
    """How many devices of type ``kind`` this machine has: 0 where it cannot start.

    Counting makes no device: ``tack.device(f"cuda:{n}")`` for ``n`` below
    the count does.
    """
    if kind not in _CONSTRUCTORS:
        available = ", ".join(sorted(_CONSTRUCTORS))
        raise ValueError(f"Unknown device type: '{kind}'. Available: {available}")
    try:
        return _backend_class(kind).device_count()
    except (ImportError, RuntimeError, OSError):
        return 0

tack.set_default_device

set_default_device(spec='cpu', **options) -> Device

Make spec's device the default: where arrays are allocated, and where kernels without arrays run. Returns it.

The device is tack.device(spec, **options): without options the same one as before, so arrays and compiled kernels survive another call.

Set TACK_NO_REINIT=1 to leave the default alone once one is set (useful when embedded in an ANARI device that shares the process).

Source code in packages/tack-core/src/tack/runtime/dispatch.py

def set_default_device(spec="cpu", **options) -> Device:
    """Make ``spec``'s device the default: where arrays are allocated, and where
    kernels without arrays run. Returns it.

    The device is ``tack.device(spec, **options)``: without options the same
    one as before, so arrays and compiled kernels survive another call.

    Set ``TACK_NO_REINIT=1`` to leave the default alone once one is set
    (useful when embedded in an ANARI device that shares the process).
    """
    global _current_backend
    if _current_backend is not None and env_flag("TACK_NO_REINIT"):
        return as_device(_current_backend)
    _current_backend = device(spec, **options)._backend
    return as_device(_current_backend)

tack.get_default_device

get_default_device() -> Device

The default device: the innermost with tack.device(...) block's, else the one set_default_device chose, else the CPU.

Source code in packages/tack-core/src/tack/runtime/dispatch.py

def get_default_device() -> Device:
    """The default device: the innermost ``with tack.device(...)`` block's, else
    the one `set_default_device` chose, else the CPU."""
    return as_device(get_backend())

tack.set_debug

set_debug(enabled: bool = True)

Check every array index in every kernel launched from now on.

A launch that went out of bounds raises IndexError naming the access; the access itself is skipped. Kernels compile separately in debug mode and run slower; each launch waits to finish. TACK_DEBUG=1 turns it on when Tack is imported.

Source code in packages/tack-core/src/tack/runtime/dispatch.py

def set_debug(enabled: bool = True):
    """Check every array index in every kernel launched from now on.

    A launch that went out of bounds raises ``IndexError`` naming the access;
    the access itself is skipped. Kernels compile separately in debug mode
    and run slower; each launch waits to finish. ``TACK_DEBUG=1`` turns it on
    when Tack is imported.
    """
    global _debug
    _debug = bool(enabled)

tack.runtime.dispatch.get_backend

get_backend()

The default device's backend: the innermost with tack.device(...) block's, else the one set_default_device chose, starting the CPU if none was.

Source code in packages/tack-core/src/tack/runtime/dispatch.py

def get_backend():
    """The default device's backend: the innermost ``with tack.device(...)``
    block's, else the one `set_default_device` chose, starting the CPU if
    none was."""
    global _current_backend
    if _open_scopes:     # a failed thread-local lookup costs ~0.2 us a launch
        stack = _scope.__dict__.get("stack")
        if stack:
            return stack[-1]
    if _current_backend is None:
        _current_backend = device("cpu")._backend
    return _current_backend

tack.sync

sync()

Wait until every kernel launched so far, on every device, has finished.

Reading or writing an array from the host (numpy, tack.copyto, indexing, reductions, DLPack export) already waits for the kernels before it, so most code never needs this. It is for timing kernels, and for memory shared with another library that reads it without Tack.

Source code in packages/tack-core/src/tack/runtime/dispatch.py

def sync():
    """Wait until every kernel launched so far, on every device, has finished.

    Reading or writing an array from the host (``numpy``, ``tack.copyto``,
    indexing, reductions, DLPack export) already waits for the kernels
    before it, so most code never needs this. It is for timing kernels, and
    for memory shared with another library that reads it without Tack.
    """
    for backend in live_backends():
        backend.synchronize()

tack.sync_async async

sync_async()

sync for asyncio: await tack.sync_async() waits for every kernel launched so far, on every device, without blocking the event loop.

On CUDA, HIP, Metal and Level Zero the wait runs in a worker thread and the loop goes on; launches made after the call are not waited for. A launch's error is raised here. On the CPU, whose launches complete before returning, it returns at once. See tack.runtime.aio.

Source code in packages/tack-core/src/tack/runtime/dispatch.py

async def sync_async():
    """`sync` for ``asyncio``: ``await tack.sync_async()`` waits for every
    kernel launched so far, on every device, without blocking the event
    loop.

    On CUDA, HIP, Metal and Level Zero the wait runs in a worker thread and
    the loop goes on; launches made after the call are not waited for. A
    launch's error is raised here. On the CPU, whose launches complete
    before returning, it returns at once. See `tack.runtime.aio`.
    """
    from tack.runtime.aio import sync_async as _sync_async
    await _sync_async()

Arrays

tack.Array holds storage on one device (a.device), chosen at allocation. Its shape counts elements; its dtype describes one scalar, vector or matrix element. tack.dataset.Field instead couples coefficient values with a space on a topology. See Arrays and Types for examples.

tack.Array

Array(
    dtype: ScalarType,
    shape: tuple[int, ...],
    buffer: DeviceBuffer,
    writable: bool = True,
    device=None,
)

An n-dimensional array of scalar, vector or matrix elements on a device.

shape counts elements and dtype describes one element, so an array of shape (n,) and dtype vec3(f32) holds 3 * n scalars and downloads as (n, 3).

layout says how a vector or matrix array stores its components: "aos" (the default) interleaves them in one buffer, "soa" keeps one scalar array per component (SoaArray). An array of scalars says "aos".

Source code in packages/tack-core/src/tack/lang/array.py

def __init__(self, dtype: ScalarType, shape: tuple[int, ...], buffer: DeviceBuffer,
             writable: bool = True, device=None):
    self._dtype = dtype
    self._shape = tuple(shape)
    self._buffer = buffer
    self._writable = writable
    # The backend that made the buffer. Every Tack constructor says; an
    # array built without one takes the default device when asked.
    self._device = device

device property

device

The tack.Device this array's memory belongs to; kernels given the array run there.

scalar_dtype property

scalar_dtype

The scalar component type used by backend storage and code generation.

dtype property

dtype

The type of one element, including vector or matrix dimensions.

shape property

shape

Array dimensions, excluding an element's component dimensions.

ndim property

ndim: int

Number of array dimensions, excluding an element's components.

storage_shape property

storage_shape

The scalar storage dimensions (components are flattened).

scalar_size property

scalar_size

Number of scalar components in storage.

nbytes property

nbytes: int

Bytes of storage the elements occupy.

size property

size: int

Total number of elements in the array.

to_device

to_device(device) -> Array

This array on device (a tack.Device or a spec such as "cpu"): itself if its memory is already usable there, else a copy.

Source code in packages/tack-core/src/tack/lang/array.py

def to_device(self, device) -> 'Array':
    """This array on ``device`` (a `tack.Device` or a spec such as ``"cpu"``):
    itself if its memory is already usable there, else a copy."""
    from tack.runtime.dispatch import backend_of
    target = backend_of(device)
    if target is self._backend or target.memory_domain == self._backend.memory_domain:
        return self
    return array(self.numpy(), self.dtype, device=target)

numpy

numpy() -> np.ndarray

Copy the values into a new NumPy array.

Vector and matrix components become trailing NumPy dimensions: an array of shape (n,) and dtype vec3(f32) gives (n, 3).

Source code in packages/tack-core/src/tack/lang/array.py

def numpy(self) -> np.ndarray:
    """Copy the values into a new NumPy array.

    Vector and matrix components become trailing NumPy dimensions: an
    array of shape ``(n,)`` and dtype ``vec3(f32)`` gives ``(n, 3)``.
    """
    arr = self._buffer.to_numpy()
    # The buffer keeps the shape it was allocated with, so a reshaped
    # view has to impose its own -- the same data either way.
    if arr.shape != self._host_shape:
        arr = arr.reshape(self._host_shape)
    return arr

__array__

__array__(dtype=None, copy=None)

np.asarray(a): the values as a NumPy array (always a copy).

Source code in packages/tack-core/src/tack/lang/array.py

def __array__(self, dtype=None, copy=None):
    """``np.asarray(a)``: the values as a NumPy array (always a copy)."""
    if copy is False:
        raise ValueError("a Tack array cannot be viewed by NumPy without a copy")
    arr = self.numpy()
    return arr if dtype is None else arr.astype(dtype, copy=False)

numpy_async async

numpy_async() -> np.ndarray

numpy for asyncio: await a.numpy_async().

Where reading waits for queued launches (CUDA, HIP, Metal, Level Zero), the wait and the copy run in a worker thread and the event loop goes on; elsewhere it is numpy. The result holds every write queued before the call; one queued later by another task may or may not be in it.

Source code in packages/tack-core/src/tack/lang/array.py

async def numpy_async(self) -> np.ndarray:
    """`numpy` for ``asyncio``: ``await a.numpy_async()``.

    Where reading waits for queued launches (CUDA, HIP, Metal, Level
    Zero), the wait and the copy run in a worker thread and the event
    loop goes on; elsewhere it is `numpy`. The result holds every write
    queued before the call; one queued later by another task may or may
    not be in it.
    """
    from tack.runtime.aio import numpy_async
    return await numpy_async(self)

fill

fill(value)

Set every element to value.

A vector or matrix array also takes one element's value, a sequence of its components: colors.fill([1.0, 1.0, 1.0]).

Source code in packages/tack-core/src/tack/lang/array.py

def fill(self, value):
    """Set every element to ``value``.

    A vector or matrix array also takes one element's value, a sequence
    of its components: ``colors.fill([1.0, 1.0, 1.0])``.
    """
    self._check_writable()
    if np.ndim(value) == 0:
        if not self._backend.fill_array(self, value):
            self._buffer.fill(value)
        return
    vector_shape = self._vector_shape()
    element = np.asarray(value)
    if vector_shape is None:
        raise TypeError(
            f"fill() takes a scalar for an array of scalars, "
            f"got a value of shape {element.shape}")
    element_shape = vector_shape[len(self._logical_shape):]
    if element.shape != element_shape:
        raise ValueError(
            f"fill() value has shape {element.shape}; "
            f"the elements of this array have shape {element_shape}")
    from tack.runtime.device_fill import element_pattern
    if not self._buffer.fill_bytes(element_pattern(element, self.scalar_dtype.numpy_dtype)):
        self._upload(np.broadcast_to(element.astype(self.scalar_dtype.numpy_dtype),
                                     vector_shape))

sum

sum()

Return a float sum; parallel floating addition order may vary.

Source code in packages/tack-core/src/tack/lang/array.py

def sum(self):
    """Return a float sum; parallel floating addition order may vary."""
    return self._reduce('sum')

min

min()

Return a float minimum; propagate NaNs and prefer negative zero.

Source code in packages/tack-core/src/tack/lang/array.py

def min(self):
    """Return a float minimum; propagate NaNs and prefer negative zero."""
    return self._reduce('min')

max

max()

Return a float maximum; propagate NaNs and prefer positive zero.

Source code in packages/tack-core/src/tack/lang/array.py

def max(self):
    """Return a float maximum; propagate NaNs and prefer positive zero."""
    return self._reduce('max')

mean

mean()

Return the mean of scalar components, or NaN for an empty array.

Source code in packages/tack-core/src/tack/lang/array.py

def mean(self):
    """Return the mean of scalar components, or NaN for an empty array."""
    if self.size == 0:
        return float('nan')
    return self.sum() / self.scalar_size

export_memory

export_memory() -> ExportedMemory

Export the array's GPU memory for cross-API sharing.

Returns an ExportedMemory with backend tag and a handle suitable for importing into another API (e.g. Dawn, Vulkan, pycuda). If the array was not allocated with exportable=True, a one-time copy into exportable memory is performed automatically.

Source code in packages/tack-core/src/tack/lang/array.py

def export_memory(self) -> ExportedMemory:
    """Export the array's GPU memory for cross-API sharing.

    Returns an ExportedMemory with backend tag and a handle suitable for
    importing into another API (e.g. Dawn, Vulkan, pycuda).
    If the array was not allocated with ``exportable=True``, a one-time
    copy into exportable memory is performed automatically.
    """
    if not hasattr(self._buffer, 'export_memory'):
        raise RuntimeError(
            f"Backend {type(self._buffer).__name__} does not support memory export")
    return self._buffer.export_memory()

__dlpack__

__dlpack__(
    *,
    stream=None,
    max_version=None,
    dl_device=None,
    copy=None,
)

Export this array as a DLPack capsule for zero-copy interop.

Usage

torch_tensor = torch.from_dlpack(tack_array) cupy_array = cupy.from_dlpack(tack_array) np_array = np.from_dlpack(tack_array) # numpy 1.25+

A consumer passing max_version >= (1, 0) gets a versioned capsule, which is the only form that can carry the read-only flag. Without it the consumer must assume the memory is writable, so NumPy marks what it gets read-only and then cannot re-export it -- which breaks round trips through anything that hands the array back.

Source code in packages/tack-core/src/tack/lang/array.py

def __dlpack__(self, *, stream=None, max_version=None, dl_device=None, copy=None):
    """Export this array as a DLPack capsule for zero-copy interop.

    Usage:
        torch_tensor = torch.from_dlpack(tack_array)
        cupy_array = cupy.from_dlpack(tack_array)
        np_array = np.from_dlpack(tack_array)  # numpy 1.25+

    A consumer passing max_version >= (1, 0) gets a versioned capsule,
    which is the only form that can carry the read-only flag. Without
    it the consumer must assume the memory is writable, so NumPy marks
    what it gets read-only and then cannot re-export it -- which breaks
    round trips through anything that hands the array back.
    """
    from tack.lang.dlpack import array_to_dlpack
    if copy is True:
        raise BufferError("Tack DLPack export does not support copy=True")
    versioned = max_version is not None and max_version[0] >= 1
    return array_to_dlpack(self, versioned=versioned)

__dlpack_device__

__dlpack_device__()

Return (device_type, device_id) for the DLPack protocol.

Source code in packages/tack-core/src/tack/lang/array.py

def __dlpack_device__(self):
    """Return (device_type, device_id) for the DLPack protocol."""
    from tack.lang.dlpack import dlpack_device
    return dlpack_device(self)

__len__

__len__() -> int

Number of elements along the first dimension.

Source code in packages/tack-core/src/tack/lang/array.py

def __len__(self) -> int:
    """Number of elements along the first dimension."""
    if not self.shape:
        raise TypeError("len() of a zero-dimensional array")
    return self.shape[0]

__getitem__

__getitem__(index)

Read one element from the host: a[i], a[i, j], or a[()] for a zero-dimensional array.

An array of scalars gives a Python number; a vector or matrix array gives a NumPy array of the element. Negative indices count from the end, as in Python. This is for inspection and for host-side decisions: on backends without host-visible memory each read is a transfer of that element, so many elements are read with numpy(). Arrays are not written this way; see tack.copyto and fill.

Source code in packages/tack-core/src/tack/lang/array.py

def __getitem__(self, index):
    """Read one element from the host: ``a[i]``, ``a[i, j]``, or ``a[()]`` for a
    zero-dimensional array.

    An array of scalars gives a Python number; a vector or matrix array
    gives a NumPy array of the element. Negative indices count from
    the end, as in Python. This is for inspection and for host-side
    decisions: on backends without host-visible memory each read is a
    transfer of that element, so many elements are read with
    ``numpy()``. Arrays are not written this way; see ``tack.copyto``
    and ``fill``.
    """
    shape = self.shape
    indices = index if isinstance(index, tuple) else (index,)
    if len(indices) != len(shape) or not all(
            isinstance(i, (int, np.integer)) and not isinstance(i, bool) for i in indices):
        if not shape:
            wanted = "a[()]"
        elif len(shape) == 1:
            wanted = "one integer index"
        else:
            wanted = f"{len(shape)} integer indices"
        raise TypeError(
            f"an array of shape {shape} takes {wanted} on the host, not {index!r}; "
            f"to read many elements use numpy()")
    flat = 0
    for i, extent in zip(indices, shape, strict=True):
        i = int(i)
        if not -extent <= i < extent:
            raise IndexError(f"index {index!r} is out of range for an array of shape {shape}")
        flat = flat * extent + (i + extent if i < 0 else i)
    vector_shape = self._vector_shape()
    if vector_shape is None:
        return self._buffer.read_range(flat, 1)[0].item()
    element_shape = vector_shape[len(shape):]
    count = int(math.prod(element_shape))
    return self._buffer.read_range(flat * count, count).reshape(element_shape)

copy

copy() -> Array

Copy all values into independent storage on the same device.

Source code in packages/tack-core/src/tack/lang/array.py

def copy(self) -> 'Array':
    """Copy all values into independent storage on the same device."""
    from tack.algorithms.copy import copy as _copy
    out = empty_like(self)
    _copy(self._scalar_view(), out._scalar_view(), self.scalar_size)
    return out

astype

astype(dtype) -> Array

Copy with a new component type or a matching vector/matrix dtype.

A scalar dtype changes the component type while preserving the element dimensions. An element dtype must have those same dimensions.

Source code in packages/tack-core/src/tack/lang/array.py

def astype(self, dtype) -> 'Array':
    """Copy with a new component type or a matching vector/matrix dtype.

    A scalar dtype changes the component type while preserving the element
    dimensions. An element dtype must have those same dimensions.
    """
    dtype = as_dtype(dtype)
    element_shape = self._vector_shape()
    if isinstance(dtype, ElementType):
        if element_shape is None or dtype.shape != element_shape[len(self._logical_shape):]:
            raise TypeError("astype cannot change an element's component dimensions")
    elif element_shape is not None:
        dtype = _element(dtype, self.dtype.shape)
    out = empty(self.shape, dtype, device=self._backend)
    from tack.algorithms.copy import copy as _copy
    _copy(self._scalar_view(), out._scalar_view(), self.scalar_size)
    return out

reshape

reshape(*shape) -> Array

Share storage under a new shape, preserving the element type.

Takes the shape as a tuple or as separate integers, as NumPy does; one dimension may be -1, inferred from the size.

Source code in packages/tack-core/src/tack/lang/array.py

def reshape(self, *shape) -> 'Array':
    """Share storage under a new shape, preserving the element type.

    Takes the shape as a tuple or as separate integers, as NumPy does;
    one dimension may be ``-1``, inferred from the size.
    """
    if len(shape) == 1 and not isinstance(shape[0], (int, np.integer)):
        shape = shape[0]
    new_shape = list(_array_shape(shape, allow_infer=True))
    if new_shape.count(-1) > 1:
        raise ValueError("reshape can infer at most one dimension")
    if -1 in new_shape:
        known = int(math.prod([n for n in new_shape if n != -1]))
        if known == 0 or self.size % known:
            raise ValueError(f"cannot reshape {self.shape} ({self.size} elements) "
                             f"to {tuple(shape)}")
        new_shape[new_shape.index(-1)] = self.size // known
    new_shape = tuple(new_shape)
    new_size = int(math.prod(new_shape))
    if new_size != self.size:
        raise ValueError(
            f"Cannot reshape {self.shape} ({self.size} elements) "
            f"to {new_shape} ({new_size} elements)")
    reshaped = _array_view(self.dtype, new_shape, self._buffer, self._writable,
                           self._device)
    hold = getattr(self, '_dlpack_hold', None)
    if hold is not None:
        reshaped._dlpack_hold = hold
    return reshaped

tack.lang.array.DeviceBuffer

Abstract interface for backend-specific array storage.

address property

address: int

Address used by kernels, for validating imported atomic storage.

host_address property

host_address: int

The address a CPU kernel reads and writes through: only memory in the host's address space has one. A CPU kernel bound to another device's memory wrote device pointers from the host and segfaulted; this raises instead.

read_range

read_range(start: int, count: int) -> np.ndarray

Elements [start, start + count) of the flat storage, as a copy.

Serves array[i] on the host. Every shipped buffer overrides it: CPU and Metal slice their host-visible views, and CUDA, HIP and Level Zero copy just the range. This default, for a buffer that does not, copies the whole buffer, which is correct and slow.

Source code in packages/tack-core/src/tack/lang/array.py

def read_range(self, start: int, count: int) -> np.ndarray:
    """Elements ``[start, start + count)`` of the flat storage, as a copy.

    Serves ``array[i]`` on the host. Every shipped buffer overrides it:
    CPU and Metal slice their host-visible views, and CUDA, HIP and
    Level Zero copy just the range. This default, for a buffer that
    does not, copies the whole buffer, which is correct and slow.
    """
    return self.to_numpy().reshape(-1)[start:start + count].copy()

fill_bytes

fill_bytes(pattern: bytes) -> bool

Fill the storage with pattern, one element's bytes, repeated.

Returns false when this buffer cannot, and the caller uploads the array instead: right for host-visible storage, which this default serves. CUDA, HIP and Level Zero override it to fill on the device (runtime/device_fill.py).

Source code in packages/tack-core/src/tack/lang/array.py

def fill_bytes(self, pattern: bytes) -> bool:
    """Fill the storage with ``pattern``, one element's bytes, repeated.

    Returns false when this buffer cannot, and the caller uploads the
    array instead: right for host-visible storage, which this default
    serves. CUDA, HIP and Level Zero override it to fill on the device
    (``runtime/device_fill.py``).
    """
    return False

synchronize

synchronize()

Wait for the launches queued on this storage's device, if any are: a backend that queues launches overrides it.

Source code in packages/tack-core/src/tack/lang/array.py

def synchronize(self):
    """Wait for the launches queued on this storage's device, if any
    are: a backend that queues launches overrides it."""

thread_binding

thread_binding()

A context manager that makes this storage usable from a thread other than the one that made it -- CUDA's context made current -- for Array.numpy_async's worker. Nothing by default.

Source code in packages/tack-core/src/tack/lang/array.py

def thread_binding(self):
    """A context manager that makes this storage usable from a thread
    other than the one that made it -- CUDA's context made current --
    for `Array.numpy_async`'s worker. Nothing by default."""
    import contextlib
    return contextlib.nullcontext()

share

share()

Hand this storage outside Tack: wait for queued launches, and have later launches that use it complete before returning. A backend that queues launches overrides it.

Source code in packages/tack-core/src/tack/lang/array.py

def share(self):
    """Hand this storage outside Tack: wait for queued launches, and
    have later launches that use it complete before returning. A
    backend that queues launches overrides it."""
    self.synchronize()

Creating arrays

tack.empty

empty(
    shape,
    dtype=f32,
    *,
    device=None,
    exportable=False,
    layout="aos",
) -> Array

A new array of uninitialized scalar, vector or matrix elements.

shape counts elements; dtype describes one element. A vector array of shape (n,) and dtype vec3(f32) downloads as (n, 3). Initialize its values with a kernel, fill or tack.copyto before reading. device is a tack.Device or a spec such as "cuda"; by default the default device. layout="soa" stores a vector or matrix array's components in one array each (see SoaArray).

Source code in packages/tack-core/src/tack/lang/array.py

def empty(shape, dtype=f32, *, device=None, exportable=False, layout="aos") -> Array:
    """A new array of uninitialized scalar, vector or matrix elements.

    ``shape`` counts elements; ``dtype`` describes one element. A vector
    array of shape ``(n,)`` and dtype ``vec3(f32)`` downloads as ``(n, 3)``.
    Initialize its values with a kernel, ``fill`` or ``tack.copyto`` before
    reading. ``device`` is a `tack.Device` or a spec such as ``"cuda"``; by
    default the default device. ``layout="soa"`` stores a vector or matrix
    array's components in one array each (see `SoaArray`).
    """
    backend = _target(device)
    shape = _array_shape(shape)
    dtype = as_dtype(dtype)
    if _check_layout(layout, dtype):
        return SoaArray(tuple(empty(shape, dtype.scalar_dtype, device=backend,
                                    exportable=exportable) for _ in range(dtype.size)),
                        dtype.shape)
    scalar = dtype.scalar_dtype if isinstance(dtype, ElementType) else dtype
    storage = (int(math.prod(shape)) * dtype.size,) if isinstance(dtype, ElementType) else shape
    buffer = _allocate(backend, scalar, storage, exportable=exportable)
    return _array_view(dtype, shape, buffer, device=backend)

tack.zeros

zeros(
    shape, dtype=f32, *, device=None, layout="aos"
) -> Array

A new array of zeros, filled on the device.

Source code in packages/tack-core/src/tack/lang/array.py

def zeros(shape, dtype=f32, *, device=None, layout="aos") -> Array:
    """A new array of zeros, filled on the device."""
    a = empty(shape, dtype, device=device, layout=layout)
    if not a._starts_zeroed:
        a.fill(0)
    return a

tack.ones

ones(
    shape, dtype=f32, *, device=None, layout="aos"
) -> Array

A new array of ones.

Source code in packages/tack-core/src/tack/lang/array.py

def ones(shape, dtype=f32, *, device=None, layout="aos") -> Array:
    """A new array of ones."""
    a = empty(shape, dtype, device=device, layout=layout)
    a.fill(1)
    return a

tack.full

full(
    shape,
    fill_value,
    dtype=None,
    *,
    device=None,
    layout="aos",
) -> Array

A new array with every element fill_value.

Without dtype, a Python float gives f32 and an int i32 (i64 if it needs it). A sequence fills vector or matrix elements; it needs an element dtype.

Source code in packages/tack-core/src/tack/lang/array.py

def full(shape, fill_value, dtype=None, *, device=None, layout="aos") -> Array:
    """A new array with every element ``fill_value``.

    Without ``dtype``, a Python float gives f32 and an int i32 (i64 if it
    needs it). A sequence fills vector or matrix elements; it needs an
    element ``dtype``.
    """
    if dtype is None:
        dtype = _host_dtype(np.asarray(fill_value), fill_value)
    a = empty(shape, dtype, device=device, layout=layout)
    a.fill(fill_value)
    return a

tack.empty_like

empty_like(
    prototype, dtype=None, *, device=None, layout=None
) -> Array

A new uninitialized array of prototype's shape and dtype.

prototype is a Tack or NumPy array; nothing is read from it. A scalar dtype changes the component type of a vector or matrix prototype and keeps its element dimensions. A Tack prototype's device and layout are the defaults.

Source code in packages/tack-core/src/tack/lang/array.py

def empty_like(prototype, dtype=None, *, device=None, layout=None) -> Array:
    """A new uninitialized array of ``prototype``'s shape and dtype.

    ``prototype`` is a Tack or NumPy array; nothing is read from it. A scalar
    ``dtype`` changes the component type of a vector or matrix prototype and
    keeps its element dimensions. A Tack prototype's device and layout are
    the defaults.
    """
    return _like(prototype, dtype, device, layout)

tack.zeros_like

zeros_like(
    prototype, dtype=None, *, device=None, layout=None
) -> Array

A new array of zeros of prototype's shape and dtype.

Source code in packages/tack-core/src/tack/lang/array.py

def zeros_like(prototype, dtype=None, *, device=None, layout=None) -> Array:
    """A new array of zeros of ``prototype``'s shape and dtype."""
    a = _like(prototype, dtype, device, layout)
    if not a._starts_zeroed:
        a.fill(0)
    return a

tack.ones_like

ones_like(
    prototype, dtype=None, *, device=None, layout=None
) -> Array

A new array of ones of prototype's shape and dtype.

Source code in packages/tack-core/src/tack/lang/array.py

def ones_like(prototype, dtype=None, *, device=None, layout=None) -> Array:
    """A new array of ones of ``prototype``'s shape and dtype."""
    a = _like(prototype, dtype, device, layout)
    a.fill(1)
    return a

tack.full_like

full_like(
    prototype,
    fill_value,
    dtype=None,
    *,
    device=None,
    layout=None,
) -> Array

A new array of prototype's shape and dtype, every element fill_value.

Source code in packages/tack-core/src/tack/lang/array.py

def full_like(prototype, fill_value, dtype=None, *, device=None, layout=None) -> Array:
    """A new array of ``prototype``'s shape and dtype, every element ``fill_value``."""
    a = _like(prototype, dtype, device, layout)
    a.fill(fill_value)
    return a

tack.array

array(
    obj, dtype=None, *, device=None, layout=None
) -> Array

A new array holding a copy of obj: a NumPy array, a sequence or a Tack array.

Without dtype, all of obj's dimensions are array dimensions, a NumPy array keeps its dtype, and Python floats and ints become f32 and i32 (i64 if a value needs it). With an element dtype such as vec3(f32), the trailing dimensions are the elements' components. A Tack array is copied on its device unless device names another. layout is "aos" or "soa" for a vector or matrix array (see SoaArray); by default a Tack array keeps its own and other data is interleaved.

Source code in packages/tack-core/src/tack/lang/array.py

def array(obj, dtype=None, *, device=None, layout=None) -> Array:
    """A new array holding a copy of ``obj``: a NumPy array, a sequence or a Tack array.

    Without ``dtype``, all of ``obj``'s dimensions are array dimensions, a
    NumPy array keeps its dtype, and Python floats and ints become f32 and
    i32 (i64 if a value needs it). With an element dtype such as
    ``vec3(f32)``, the trailing dimensions are the elements' components.
    A Tack array is copied on its device unless ``device`` names another.
    ``layout`` is ``"aos"`` or ``"soa"`` for a vector or matrix array (see
    `SoaArray`); by default a Tack array keeps its own and other data is
    interleaved.
    """
    if isinstance(obj, Array):
        if layout is not None and layout != obj.layout:
            if device is not None:
                obj = obj.to_device(device)
            out = empty(obj.shape, obj.dtype if dtype is None else dtype,
                        device=obj._backend, layout=layout)
            copyto(out, obj)
            return out
        if device is not None and _target(device).memory_domain != obj._backend.memory_domain:
            moved = obj.to_device(device)
            return moved if dtype is None else moved.astype(dtype)
        return obj.copy() if dtype is None else obj.astype(dtype)
    host = np.asarray(obj)
    dtype = _host_dtype(host, obj) if dtype is None else as_dtype(dtype)
    shape = host.shape
    if isinstance(dtype, ElementType):
        rank = len(dtype.shape)
        if len(shape) < rank or shape[-rank:] != dtype.shape:
            raise ValueError(f"input shape {shape} must end with element shape {dtype.shape}")
        shape = shape[:-rank]
    out = empty(shape, dtype, device=device, layout=layout or "aos")
    out._upload(host.astype(out.scalar_dtype.numpy_dtype, copy=False))
    return out

tack.asarray

asarray(
    obj, dtype=None, *, device=None, layout=None
) -> Array

obj itself if it is a Tack array of dtype (or any, without one) on device (or any, without one) in layout (or either, without one), else array(obj, dtype, device=device, layout=layout).

Source code in packages/tack-core/src/tack/lang/array.py

def asarray(obj, dtype=None, *, device=None, layout=None) -> Array:
    """``obj`` itself if it is a Tack array of ``dtype`` (or any, without one)
    on ``device`` (or any, without one) in ``layout`` (or either, without one),
    else `array(obj, dtype, device=device, layout=layout)`."""
    if isinstance(obj, Array):
        if layout is not None and layout != obj.layout:
            return array(obj, dtype, device=device, layout=layout)
        if device is not None:
            obj = obj.to_device(device)
        if dtype is None or as_dtype(dtype) == obj.dtype:
            return obj
        return obj.astype(dtype)
    return array(obj, dtype, device=device, layout=layout)

tack.arange

arange(
    start, stop=None, step=1, dtype=None, *, device=None
) -> Array

Evenly spaced values in [start, stop), as np.arange.

Without dtype, integer arguments give i32 (i64 if a value needs it) and any float argument gives f32.

Source code in packages/tack-core/src/tack/lang/array.py

def arange(start, stop=None, step=1, dtype=None, *, device=None) -> Array:
    """Evenly spaced values in ``[start, stop)``, as ``np.arange``.

    Without ``dtype``, integer arguments give i32 (i64 if a value needs it)
    and any float argument gives f32.
    """
    if stop is None:
        start, stop = 0, start
    values = np.arange(start, stop, step)
    if dtype is None:
        dtype = _host_dtype(values, [start, stop, step])
        if values.size == 0 and values.dtype.kind in 'iu':
            dtype = i32 if all(integer_type_for_value(int(v)) is i32 for v in (start, stop)) else i64
    return array(values, dtype, device=device)

tack.concat

concat(arrays) -> Array

Join a sequence of arrays end to end into a new one-dimensional array.

All arrays must have the same dtype and be on one device, where the copy runs.

Source code in packages/tack-core/src/tack/lang/array.py

def concat(arrays) -> Array:
    """Join a sequence of arrays end to end into a new one-dimensional array.

    All arrays must have the same dtype and be on one device, where the copy
    runs.
    """
    arrays = list(arrays)
    if not arrays:
        raise ValueError("concat requires at least one array")
    dt = arrays[0].dtype
    for a in arrays[1:]:
        if a.dtype != dt:
            raise TypeError(
                f"concat: all arrays must have the same dtype, got {dt} and {a.dtype}")
    from tack.runtime.dispatch import same_device
    same_device("concat()", arrays)
    total = sum(a.size for a in arrays)
    result = empty((total,), dt, device=arrays[0]._backend)
    from tack.algorithms.copy import copy_with_offset
    offset = 0
    for a in arrays:
        copy_with_offset(a, result, offset, a.size)
        offset += a.size
    return result

Structure of arrays

A vector or matrix array stored one scalar array per component (layout="soa"); kernels index it as any vector array.

tack.soa

soa(*components, dtype=None) -> SoaArray

A vector array whose components are the given arrays, without copying.

tack.soa(x, y, z) is an array of vec3 elements whose element i is (x[i], y[i], z[i]): the structure-of-arrays layout, for data that arrives that way, such as a simulation's coordinate arrays or another library's tensors (from_dlpack, array_from_ptr). Kernels read and write it as they do any vector array, and writes go to the component arrays. The components are arrays of scalars of one dtype and one shape, on one device. An element dtype of the same component type makes a matrix array instead, components row-major: tack.soa(a, b, c, d, dtype=tack.mat2()).

Source code in packages/tack-core/src/tack/lang/array.py

def soa(*components, dtype=None) -> SoaArray:
    """A vector array whose components are the given arrays, without copying.

    ``tack.soa(x, y, z)`` is an array of ``vec3`` elements whose element
    ``i`` is ``(x[i], y[i], z[i])``: the structure-of-arrays layout, for
    data that arrives that way, such as a simulation's coordinate arrays
    or another library's tensors (``from_dlpack``, ``array_from_ptr``).
    Kernels read and write it as they do any vector array, and writes go
    to the component arrays. The components are arrays of scalars of one
    dtype and one shape, on one device. An element ``dtype`` of the same
    component type makes a matrix array instead, components row-major:
    ``tack.soa(a, b, c, d, dtype=tack.mat2())``.
    """
    if len(components) == 1 and isinstance(components[0], (list, tuple)):
        components = tuple(components[0])
    if not components:
        raise TypeError("soa() takes the component arrays, at least one")
    for c in components:
        if not isinstance(c, Array):
            raise TypeError(f"soa() takes Tack arrays, not {type(c).__name__}; make one "
                            f"with tack.array(), tack.from_dlpack() or tack.array_from_ptr()")
        if hasattr(c, '_vector_n'):
            raise TypeError(f"soa() takes arrays of scalars, one per component, "
                            f"not an array of {c.dtype}")
    first = components[0]
    for c in components[1:]:
        if c.scalar_dtype is not first.scalar_dtype or c.shape != first.shape:
            raise ValueError(
                f"soa() components must share a dtype and a shape; got {first.dtype} "
                f"{first.shape} and {c.dtype} {c.shape}")
    from tack.runtime.dispatch import same_device
    same_device("soa()", components)
    if dtype is None:
        element = vec(len(components), first.scalar_dtype)
    else:
        element = as_dtype(dtype)
        if not isinstance(element, ElementType) or element.size != len(components):
            raise TypeError(f"soa() was given {len(components)} components; dtype must be "
                            f"a vector or matrix of {len(components)}, not {element}")
        if element.scalar_dtype is not first.scalar_dtype:
            raise TypeError(f"soa() components are {first.scalar_dtype}; dtype {element} "
                            f"has {element.scalar_dtype} components")
    return SoaArray(components, element.shape)

tack.SoaArray

SoaArray(components, element_shape)

Bases: Array

A vector or matrix array stored as one scalar array per component.

Made by soa from arrays that exist already (from another library, through from_dlpack or array_from_ptr), or by a constructor given layout="soa". It is an Array like any other: a kernel reads a[i] as a vector, stores, updates and atomically updates it, and only the generated code differs, reading component c of element i from a.components[c][i] rather than from i * n + c of one buffer. A kernel compiled for one layout is not used for the other.

With no single buffer, an SOA array is not exported whole (DLPack, export_memory); its components are.

Source code in packages/tack-core/src/tack/lang/array.py

def __init__(self, components, element_shape):
    first = components[0]
    super().__init__(first.scalar_dtype, (len(components) * first.size,), None,
                     all(c._writable for c in components), first._backend)
    self._components = tuple(components)
    self._vector_n = len(components)
    self._logical_shape = first.shape
    if len(element_shape) == 2:
        self._matrix_shape = tuple(element_shape)

components property

components: tuple

The scalar arrays holding each component, in order (row-major for a matrix): the arrays soa was given, or the ones a constructor made.

numpy

numpy() -> np.ndarray

Copy the values into a new NumPy array, components as the trailing dimensions, exactly as for an interleaved array.

Source code in packages/tack-core/src/tack/lang/array.py

def numpy(self) -> np.ndarray:
    """Copy the values into a new NumPy array, components as the trailing
    dimensions, exactly as for an interleaved array."""
    return np.stack([c.numpy() for c in self._components], axis=-1).reshape(
        self._host_shape)

Copying values

tack.copyto

copyto(dst: Array, src)

Copy src into the existing array dst, as np.copyto does.

src is host data -- a NumPy array, a sequence or a scalar -- whose shape is dst.numpy()'s or broadcasts to it, converted to dst's dtype; or a Tack array of dst's shape on dst's device, copied there (src.to_device(dst.device) brings one from another device).

It only writes into Tack arrays. numpy reads one into a new NumPy array; np.copyto copies between NumPy arrays.

Source code in packages/tack-core/src/tack/lang/array.py

def copyto(dst: Array, src):
    """Copy ``src`` into the existing array ``dst``, as ``np.copyto`` does.

    ``src`` is host data -- a NumPy array, a sequence or a scalar -- whose
    shape is ``dst.numpy()``'s or broadcasts to it, converted to ``dst``'s
    dtype; or a Tack array of ``dst``'s shape on ``dst``'s device, copied
    there (``src.to_device(dst.device)`` brings one from another device).

    It only writes into Tack arrays. `numpy` reads one into a new NumPy
    array; ``np.copyto`` copies between NumPy arrays.
    """
    if not isinstance(dst, Array):
        raise TypeError(f"copyto() writes into a Tack array, not {type(dst).__name__}")
    if isinstance(src, Array):
        if src.shape != dst.shape or src.scalar_size != dst.scalar_size:
            raise ValueError(f"copyto() from an array of shape {src.shape} and dtype "
                             f"{src.dtype} into one of shape {dst.shape} and dtype {dst.dtype}")
        from tack.runtime.dispatch import same_device
        same_device("copyto()", (dst, src))
        dst._check_writable()
        from tack.algorithms.copy import copy as _copy
        if src.layout == dst.layout == "aos":
            _copy(src._scalar_view(), dst._scalar_view(), dst.scalar_size)
            return
        # An SOA array has no one buffer: copy element by element, which
        # reads and writes either layout.
        if getattr(src.dtype, 'shape', None) != getattr(dst.dtype, 'shape', None):
            raise ValueError(f"copyto() from an array of {src.dtype} elements into one "
                             f"of {dst.dtype} elements")
        _copy(src.reshape(src.size), dst.reshape(dst.size), dst.size)
        return
    host = np.asarray(src, dtype=dst.scalar_dtype.numpy_dtype)
    expected = dst._host_shape
    if host.shape != expected:
        try:
            host = np.broadcast_to(host, expected)
        except ValueError:
            raise ValueError(f"copyto() into an array whose values have shape {expected}: "
                             f"cannot broadcast shape {host.shape}") from None
    dst._upload(host)

Implicit arrays

Arrays that compute each value from its index, in tack.arrays; a kernel indexes any of them as it does a stored array (a[k], len(a)).

tack.arrays.CountingArray

CountingArray(size, start=0, step=1, dtype=None)

Bases: _Implicit

start + step * index for index < size (VTK's affine array): integers are dtype, by default i32 unless a value needs i64 (ids past 2**31 - 1), else f32.

Source code in packages/tack-core/src/tack/arrays.py

def __init__(self, size, start=0, step=1, dtype=None):
    self.start = start
    self.step = step
    integral = all(isinstance(v, (int, np.integer)) for v in (start, step))
    if dtype is None and integral:
        reach = max(abs(start), abs(start + step * max(int(size) - 1, 0)))
        dtype = tack.i64 if reach > 2**31 - 1 else tack.i32
    _metadata(self, size, None, dtype or tack.f32)

tack.arrays.ConstantArray

ConstantArray(value, size, dtype=None, device=None)

Bases: _Implicit

size copies of value: a scalar, or a vector (a sequence of components), as VTK's constant arrays of several components. dtype defaults to i32 for integers and f32 otherwise; the value is kept in a one-element field of it, so it is exactly that type in the kernel, on device (by default the default device).

Source code in packages/tack-core/src/tack/arrays.py

def __init__(self, value, size, dtype=None, device=None):
    host = np.asarray(value)
    if dtype is None:
        dtype = tack.i32 if host.dtype.kind in "iu" else tack.f32
    self.value = _store(host.reshape(1, -1) if host.ndim else host.reshape(1), dtype,
                        device)
    _metadata(self, size, host.shape[0] if host.ndim else None, dtype)

tack.arrays.CartesianProduct

CartesianProduct(
    x, y=(0.0,), z=(0.0,), dtype=tack.f32, device=None
)

Bases: _Implicit

The points of a grid whose lines run along the axes: (x[i], y[j], z[k]) at index i + nx * (j + ny * k). Only the three axes are stored.

Source code in packages/tack-core/src/tack/arrays.py

def __init__(self, x, y=(0.0,), z=(0.0,), dtype=tack.f32, device=None):
    if device is None:      # host axes go where the device axes are
        device = next((a._device for a in (x, y, z) if isinstance(a, _DeviceArray)), None)
    self.xs, self.ys, self.zs = (self._axis(a, dtype, device) for a in (x, y, z))
    self.dims = (self.xs.shape[0], self.ys.shape[0], self.zs.shape[0])
    self.px, self.py = self.dims[0], self.dims[1]
    _metadata(self, math.prod(self.dims), 3, dtype)

point_dims property

point_dims

dims without trailing ones: the dimensions of the grid of cells.

tack.arrays.UniformCoordinates

UniformCoordinates(
    dims,
    origin=(0.0, 0.0, 0.0),
    spacing=(1.0, 1.0, 1.0),
    direction=None,
    dtype=tack.f32,
    device=None,
)

Bases: _Implicit

The points of an image: origin + direction @ ((i, j, k) * spacing) at index i + nx * (j + ny * k), as VTK's image data (with its direction matrix) and Viskores' uniform point coordinates. Nothing per point is stored.

dims has one to three sizes; origin and spacing as many values or three, the rest taken as 0 and 1, and direction is 3 x 3 or, for a 2D image, 2 x 2.

Source code in packages/tack-core/src/tack/arrays.py

def __init__(self, dims, origin=(0.0, 0.0, 0.0), spacing=(1.0, 1.0, 1.0), direction=None,
             dtype=tack.f32, device=None):
    dims = tuple(int(d) for d in dims)
    if not 1 <= len(dims) <= 3:
        raise ValueError(f"an image has one to three sizes, not {dims!r}")
    dims += (1,) * (3 - len(dims))
    self.dims = dims
    self.px, self.py = dims[0], dims[1]
    origin, spacing = _three(origin, 0.0, "origin"), _three(spacing, 1.0, "spacing")
    self.ox, self.oy, self.oz = origin
    self.sx, self.sy, self.sz = spacing
    matrix = np.eye(3)
    if direction is not None:
        given = np.asarray(direction, float)
        if given.shape == (2, 2):
            matrix[:2, :2] = given
        elif given.size == 9:
            matrix = given.reshape(3, 3)
        else:
            raise ValueError(f"direction is a 3 x 3 matrix, or 2 x 2 for a 2D image, not "
                             f"{given.shape}")
    if direction is not None and not np.allclose(matrix, np.eye(3)):
        self.ORIENTED = 1
    # Always present, so an f64 grid's scalars are f64 in the kernel; it
    # puts the grid on ``device`` (by default the default device).
    self.direction = _store(matrix.reshape(-1), dtype, device)
    self._origin, self._spacing, self._matrix = origin, spacing, matrix
    _metadata(self, math.prod(dims), 3, dtype)

tack.arrays.Permutation

Permutation(ids, base)

Bases: _Implicit

base[ids[k]]: base's values in the order ids (integers, explicit or not) gives, as VTK's indexed array.

Source code in packages/tack-core/src/tack/arrays.py

def __init__(self, ids, base):
    self.ids, self.base = _operands(ids, base)
    _metadata(self, size_of(self.ids), width_of(self.base), dtype_of(self.base))

tack.arrays.Concatenate

Concatenate(*arrays)

The arrays end to end, as VTK's composite array: of one width and type.

Source code in packages/tack-core/src/tack/arrays.py

def Concatenate(*arrays):
    """The arrays end to end, as VTK's composite array: of one width and type."""
    if not arrays:
        raise ValueError("Concatenate needs at least one array")
    parts = _operands(*arrays)
    for part in parts[1:]:
        if (width_of(part), dtype_of(part)) != (width_of(parts[0]), dtype_of(parts[0])):
            raise TypeError("concatenated arrays must have one width and one type")
    joined = parts[0]
    for part in parts[1:]:
        joined = _Concatenation(joined, part)
    return joined

tack.arrays.View

View(base, start, size)

Bases: _Implicit

base[start + k] for k < size: a stretch of another array.

Source code in packages/tack-core/src/tack/arrays.py

def __init__(self, base, start, size):
    self.base = _operand(base)
    self.start = int(start)
    _metadata(self, size, width_of(self.base), dtype_of(self.base))

tack.arrays.Strided

Strided(values, stride, offset=0, size=None)

Bases: _Implicit

values[offset + stride * k] of a scalar device array, as VTK's strided array: one column of interleaved data, or every n-th value.

Source code in packages/tack-core/src/tack/arrays.py

def __init__(self, values, stride, offset=0, size=None):
    values = _operand(values)
    if width_of(values):
        raise TypeError("Strided reads a scalar field; ExtractComponent reads a vector's")
    self.values = values
    self.stride, self.offset = int(stride), int(offset)
    count = -(-(size_of(values) - self.offset) // self.stride) if size is None else size
    _metadata(self, count, None, dtype_of(values))

tack.arrays.ExtractComponent

ExtractComponent(base, component)

Bases: _Implicit

Component c of a vector array's values.

Source code in packages/tack-core/src/tack/arrays.py

def __init__(self, base, component):
    self.base = _operand(base)
    if not width_of(self.base) or not 0 <= component < width_of(self.base):
        raise ValueError(f"component {component} of an array of width "
                         f"{width_of(self.base)}")
    self.COMPONENT = int(component)
    _metadata(self, size_of(self.base), None, dtype_of(self.base))

tack.arrays.Components

Components(*arrays)

Two to four scalar arrays as the components of a vector array, as Viskores' composite vector.

Source code in packages/tack-core/src/tack/arrays.py

def Components(*arrays):
    """Two to four scalar arrays as the components of a vector array, as Viskores'
    composite vector."""
    parts = _operands(*arrays)
    if not 2 <= len(parts) <= 4 or any(width_of(p) for p in parts):
        raise TypeError("Components takes two to four scalar arrays")
    if len({size_of(p) for p in parts}) != 1:
        raise ValueError("the component arrays differ in size")
    if len({dtype_of(p) for p in parts}) != 1:
        raise TypeError("the component arrays differ in type: Cast them to one")
    out = {2: _Components2, 3: _Components3, 4: _Components4}[len(parts)]()
    for c, part in enumerate(parts):
        setattr(out, f"c{c}", part)
    _metadata(out, size_of(parts[0]), len(parts), dtype_of(parts[0]))
    return out

tack.arrays.Function

Function(func, size, dtype=tack.f32, width=None)

Bases: _Implicit

func(k) for k < size: any @tack.func of the index, as VTK's vtkStdFunctionArray -- but evaluated in the kernel, on every backend. Its values are dtype (and vectors of width) as the function returns them.

Source code in packages/tack-core/src/tack/arrays.py

def __init__(self, func, size, dtype=tack.f32, width=None):
    self.func = _device_function(func)
    _metadata(self, size, width, dtype)

tack.arrays.Transform

Transform(func, base, dtype=None, width=None)

Bases: _Implicit

func(base[k]): another array's values through a @tack.func. dtype and width are the function's results, by default base's.

Source code in packages/tack-core/src/tack/arrays.py

def __init__(self, func, base, dtype=None, width=None):
    self.func = _device_function(func)
    self.base = _operand(base)
    _metadata(self, size_of(self.base), width if width is not None else width_of(self.base),
              dtype or dtype_of(self.base))

tack.arrays.Cast

Cast(base, dtype)

base's values converted to dtype, component by component.

Source code in packages/tack-core/src/tack/arrays.py

def Cast(base, dtype):
    """``base``'s values converted to ``dtype``, component by component."""
    if dtype not in _CASTS:
        raise TypeError(f"Cast converts to {', '.join(t.name for t in _CASTS)}")
    return Transform(_CASTS[dtype], base, dtype=dtype)

tack.arrays.RandomUniform

RandomUniform(size, seed=0, low=0.0, high=1.0)

Bases: _Implicit

size uniform draws in [low, high), value k a pure function of k and seed (tack.random: the same on every backend, and from tack.random.np_uniform on the host). Indices repeat past 2**32.

Source code in packages/tack-core/src/tack/arrays.py

def __init__(self, size, seed=0, low=0.0, high=1.0):
    self.seed = int(seed)
    self.low, self.span = float(low), float(high) - float(low)
    _metadata(self, size, None, tack.f32)

tack.arrays.RandomNormal

RandomNormal(size, seed=0, mean=0.0, stddev=1.0)

Bases: _Implicit

size normal draws of mean and stddev, as RandomUniform.

Source code in packages/tack-core/src/tack/arrays.py

def __init__(self, size, seed=0, mean=0.0, stddev=1.0):
    self.seed = int(seed)
    self.mean, self.stddev = float(mean), float(stddev)
    _metadata(self, size, None, tack.f32)

tack.arrays.materialize

materialize(values)

values as a device array: itself if it is one, else computed by a kernel.

Source code in packages/tack-core/src/tack/arrays.py

def materialize(values):
    """``values`` as a device array: itself if it is one, else computed by a kernel."""
    if not isinstance(values, _Implicit):
        return values
    if isinstance(values, _Explicit):
        return values.values
    out = tack.empty((values.size,), element_type(values), device=values._device)
    if values.size:
        _fill(values, out)
    return out

tack.arrays.size_of

size_of(values)

The number of values.

Source code in packages/tack-core/src/tack/arrays.py

def size_of(values):
    """The number of values."""
    return values.size

tack.arrays.width_of

width_of(values)

Components per value: 3 for 3-vectors, None for scalars.

Source code in packages/tack-core/src/tack/arrays.py

def width_of(values):
    """Components per value: 3 for 3-vectors, None for scalars."""
    if isinstance(values, _Implicit):
        return values.width
    return getattr(values, "_vector_n", None)

tack.arrays.dtype_of

dtype_of(values)

The type of each component.

Source code in packages/tack-core/src/tack/arrays.py

def dtype_of(values):
    """The type of each component."""
    return values.scalar_dtype if isinstance(values, _DeviceArray) else values.dtype

tack.arrays.to_host

to_host(values)

The values as a host array: (size,) for scalars, (size, width) for vectors.

Source code in packages/tack-core/src/tack/arrays.py

def to_host(values):
    """The values as a host array: ``(size,)`` for scalars, ``(size, width)`` for vectors."""
    return values.numpy()

Element types

tack.ElementType dataclass

ElementType(
    scalar_dtype: ScalarType, shape: tuple[int, ...]
)

A fixed-size vector or matrix element, independent of an array's shape.

Components are stored together, without padding, in row-major order. scalar_dtype is the component type; shape describes one element.

tack.vec cached

vec(n: int, dtype: ScalarType = f32) -> ElementType

The type of an n-component vector; no storage is allocated.

Source code in packages/tack-core/src/tack/lang/types.py

@lru_cache(maxsize=None, typed=True)
def vec(n: int, dtype: ScalarType = f32) -> ElementType:
    """The type of an n-component vector; no storage is allocated."""
    return ElementType(dtype, (n,))

tack.vec2

vec2(dtype: ScalarType = f32) -> ElementType

A two-component vector element type; in a kernel, tack.vec2(...) builds one.

In kernels it takes 2 scalars, one scalar for every component, or one 2-vector.

Source code in packages/tack-core/src/tack/lang/types.py

def vec2(dtype: ScalarType = f32) -> ElementType:
    """A two-component vector element type; in a kernel, ``tack.vec2(...)`` builds one.

    In kernels it takes 2 scalars, one scalar for every component, or one
    2-vector.
    """
    return vec(2, dtype)

tack.vec3

vec3(dtype: ScalarType = f32) -> ElementType

A three-component vector element type; in a kernel, tack.vec3(...) builds one.

In kernels it takes 3 scalars, one scalar for every component, or one 3-vector.

Source code in packages/tack-core/src/tack/lang/types.py

def vec3(dtype: ScalarType = f32) -> ElementType:
    """A three-component vector element type; in a kernel, ``tack.vec3(...)`` builds one.

    In kernels it takes 3 scalars, one scalar for every component, or one
    3-vector.
    """
    return vec(3, dtype)

tack.vec4

vec4(dtype: ScalarType = f32) -> ElementType

A four-component vector element type; in a kernel, tack.vec4(...) builds one.

In kernels it takes 4 scalars, one scalar for every component, or one 4-vector.

Source code in packages/tack-core/src/tack/lang/types.py

def vec4(dtype: ScalarType = f32) -> ElementType:
    """A four-component vector element type; in a kernel, ``tack.vec4(...)`` builds one.

    In kernels it takes 4 scalars, one scalar for every component, or one
    4-vector.
    """
    return vec(4, dtype)

tack.mat cached

mat(
    rows: int, columns: int, dtype: ScalarType = f32
) -> ElementType

The type of a row-major matrix element, up to 4x4.

In a kernel, tack.mat(row, ...) builds a matrix from its rows instead.

Source code in packages/tack-core/src/tack/lang/types.py

@lru_cache(maxsize=None, typed=True)
def mat(rows: int, columns: int, dtype: ScalarType = f32) -> ElementType:
    """The type of a row-major matrix element, up to 4x4.

    In a kernel, ``tack.mat(row, ...)`` builds a matrix from its rows instead.
    """
    return ElementType(dtype, (rows, columns))

tack.mat2

mat2(dtype: ScalarType = f32) -> ElementType

A 2x2 matrix element type; in a kernel, tack.mat2(...) builds one.

In kernels it takes 2 rows (lists of scalars or vectors) or 4 scalars in row-major order.

Source code in packages/tack-core/src/tack/lang/types.py

def mat2(dtype: ScalarType = f32) -> ElementType:
    """A 2x2 matrix element type; in a kernel, ``tack.mat2(...)`` builds one.

    In kernels it takes 2 rows (lists of scalars or vectors) or 4
    scalars in row-major order.
    """
    return mat(2, 2, dtype)

tack.mat3

mat3(dtype: ScalarType = f32) -> ElementType

A 3x3 matrix element type; in a kernel, tack.mat3(...) builds one.

In kernels it takes 3 rows (lists of scalars or vectors) or 9 scalars in row-major order.

Source code in packages/tack-core/src/tack/lang/types.py

def mat3(dtype: ScalarType = f32) -> ElementType:
    """A 3x3 matrix element type; in a kernel, ``tack.mat3(...)`` builds one.

    In kernels it takes 3 rows (lists of scalars or vectors) or 9
    scalars in row-major order.
    """
    return mat(3, 3, dtype)

tack.mat4

mat4(dtype: ScalarType = f32) -> ElementType

A 4x4 matrix element type; in a kernel, tack.mat4(...) builds one.

In kernels it takes 4 rows (lists of scalars or vectors) or 16 scalars in row-major order.

Source code in packages/tack-core/src/tack/lang/types.py

def mat4(dtype: ScalarType = f32) -> ElementType:
    """A 4x4 matrix element type; in a kernel, ``tack.mat4(...)`` builds one.

    In kernels it takes 4 rows (lists of scalars or vectors) or 16
    scalars in row-major order.
    """
    return mat(4, 4, dtype)

Kernel values and textures

Inside a kernel the element-type names above build values instead: tack.vec3(x, y, z), tack.vec3(s) (every component s) or tack.vec3(v) (a 3-vector), and tack.mat2([a, b], [c, d]), tack.mat2(a, b, c, d) (row-major) or tack.mat(row, ...) for any shape up to 4×4. See Vectors and Matrices.

tack.identity

identity(n)

The n-by-n identity matrix, as np.identity. Only usable inside a @tack.kernel.

Source code in packages/tack-core/src/tack/__init__.py

def identity(n):
    """The n-by-n identity matrix, as ``np.identity``. Only usable inside a @tack.kernel."""
    raise RuntimeError("identity() can only be used inside a @tack.kernel")

tack.select

select(mask, a, b)

a where mask is true, else b, component by component.

Only usable inside a kernel or device function. mask is a scalar or a comparison vector such as v > 0.0; a and b are scalars or vectors of one shape, a scalar standing for every component.

Source code in packages/tack-core/src/tack/__init__.py

def select(mask, a, b):
    """``a`` where ``mask`` is true, else ``b``, component by component.

    Only usable inside a kernel or device function. ``mask`` is a scalar
    or a comparison vector such as ``v > 0.0``; ``a`` and ``b`` are scalars
    or vectors of one shape, a scalar standing for every component.
    """
    raise RuntimeError("tack.select() can only be used inside a @tack.kernel")

tack.texture3d

texture3d(source_array, shape=None, interp='linear')

Create a 3D texture from an array for hardware-accelerated sampling.

The texture copies the array's data now. Later writes to the array -- kernel stores, from_numpy, external or DLPack writes -- are not visible to it until tex.update() copies the array again. That holds on every backend, including CPU.

Parameters:

Name	Type	Description	Default
source_array		A Tack array with f32 dtype holding WHD elements, x varying fastest.	required
shape		(W, H, D) tuple for the 3D dimensions. If the array's shape is already 3D, this can be omitted.	None
interp		Interpolation mode. Only 'linear' (the default) is implemented.	'linear'

Source code in packages/tack-core/src/tack/__init__.py

def texture3d(source_array, shape=None, interp='linear'):
    """Create a 3D texture from an array for hardware-accelerated sampling.

    The texture copies the array's data now. Later writes to the array --
    kernel stores, ``from_numpy``, external or DLPack writes -- are not
    visible to it until ``tex.update()`` copies the array again. That holds
    on every backend, including CPU.

    Args:
        source_array: A Tack array with f32 dtype holding W*H*D elements,
               x varying fastest.
        shape: (W, H, D) tuple for the 3D dimensions.  If the array's shape
               is already 3D, this can be omitted.
        interp: Interpolation mode. Only 'linear' (the default) is
               implemented.
    """
    if shape is None:
        shape = source_array.shape
    if len(shape) != 3:
        raise ValueError("texture3d requires a 3D shape")
    return Texture3D(source_array, shape_3d=shape, interp=interp)

tack.lang.array.Texture3D

Texture3D(
    source_array: Array,
    shape_3d: tuple,
    interp: str = "linear",
)

A 3D texture over a snapshot of an array, sampled trilinearly.

Created via tack.texture3d(array, shape=(W, H, D)). In kernels, tex.sample(u, v, w) samples at normalized [0,1] coordinates using trilinear interpolation.

The texture copies the array's data when it is created, and again on update(); writes to the array in between do not reach it. That holds on every backend. Where the backend samples in hardware the copy is a texture image (a CUDA or HIP array, an MTLTexture, a Level Zero image); elsewhere it is a private array the generated code interpolates in software. Either way the texture owns it, so it is released with the texture and no other texture can be handed it.

Source code in packages/tack-core/src/tack/lang/array.py

def __init__(self, source_array: Array, shape_3d: tuple, interp: str = 'linear'):
    # Every hardware path builds a single-channel 32-bit float image,
    # so f64 data would be reinterpreted rather than converted.
    if source_array.dtype is not f32:
        raise ValueError(
            f"texture3d requires an f32 array, got {source_array.dtype}; "
            f"convert it first with array.astype(tack.f32)")
    # The generated code interpolates linearly on every backend; nothing
    # implements another mode.
    if interp != 'linear':
        raise ValueError(
            f"texture3d supports interp='linear' only, got {interp!r}")
    shape_3d = tuple(int(s) for s in shape_3d)
    W, H, D = shape_3d
    if W * H * D != source_array.size:
        raise ValueError(
            f"texture3d shape {shape_3d} holds {W * H * D} elements, but "
            f"the array has {source_array.size} elements")
    self.array = source_array
    self.shape_3d = shape_3d   # (W, H, D) logical 3D shape
    self.interp = interp

    # The texture lives where its array does, and kernels sampling it run
    # there (``dispatch.launch_backend`` reads ``_device``).
    backend = self._device = source_array._backend
    if backend.texture_in_hardware(shape_3d):
        self._storage = backend.create_texture_image(shape_3d)
    else:
        self._storage = Array(f32, (source_array.size,),
                              _allocate(backend, f32, (source_array.size,)),
                              device=backend)
    self.update()

update

update()

Copy the array's current data into the texture.

Sampling after this sees the array as it is now, on every backend.

Source code in packages/tack-core/src/tack/lang/array.py

def update(self):
    """Copy the array's current data into the texture.

    Sampling after this sees the array as it is now, on every backend.
    """
    if isinstance(self._storage, Array):
        from tack.algorithms.copy import copy as _copy
        _copy(self.array.reshape((self.array.size,)), self._storage,
              self.array.size)
    else:
        self._storage.upload(self.array)

sample

sample(u, v, w)

Sample at normalized coordinates. Only usable inside @tack.kernel.

Source code in packages/tack-core/src/tack/lang/array.py

def sample(self, u, v, w):
    """Sample at normalized coordinates. Only usable inside @tack.kernel."""
    raise RuntimeError("Texture3D.sample() can only be used inside a @tack.kernel")

Sharing memory

tack.array_from_ptr

array_from_ptr(
    ptr,
    shape,
    dtype,
    *,
    device=None,
    writable: bool = False,
) -> Array

Wrap existing device memory as a Tack array without copying.

Use this for interop with external libraries (pycuda, cupy, Catalyst) or when receiving a device pointer from a simulation framework.

The array does NOT own the memory — Tack will not free it. Read-only by default; pass writable=True to enable writes.

Raises ValueError if the pointer's memory space does not match the active backend (e.g. a CPU pointer with the CUDA backend), and TypeError if a CUDA, HIP or Level Zero pointer is not an address.

Parameters:

Name	Type	Description	Default
ptr		device pointer (integer) or backend-specific buffer object. - CPU: integer address or numpy array - Metal: MTLBuffer object (from PyObjC) - CUDA/HIP: device pointer as integer - Level Zero: device pointer as integer	required
shape		the array's dimensions, counting elements	required
dtype		scalar or element type (tack.f32, tack.vec3(), np.int32, ...)	required
device		the device whose memory this is; by default the default device	None
writable	bool	if False (default), copyto() and fill() raise errors	False

Returns:

Type	Description
Array	An Array wrapping the external memory.

Source code in packages/tack-core/src/tack/lang/array.py

def array_from_ptr(ptr, shape, dtype, *, device=None, writable: bool = False) -> Array:
    """Wrap existing device memory as a Tack array without copying.

    Use this for interop with external libraries (pycuda, cupy, Catalyst)
    or when receiving a device pointer from a simulation framework.

    The array does NOT own the memory — Tack will not free it.
    Read-only by default; pass writable=True to enable writes.

    Raises ValueError if the pointer's memory space does not match the
    active backend (e.g. a CPU pointer with the CUDA backend), and
    TypeError if a CUDA, HIP or Level Zero pointer is not an address.

    Args:
        ptr: device pointer (integer) or backend-specific buffer object.
             - CPU: integer address or numpy array
             - Metal: MTLBuffer object (from PyObjC)
             - CUDA/HIP: device pointer as integer
             - Level Zero: device pointer as integer
        shape: the array's dimensions, counting elements
        dtype: scalar or element type (tack.f32, tack.vec3(), np.int32, ...)
        device: the device whose memory this is; by default the default device
        writable: if False (default), copyto() and fill() raise errors

    Returns:
        An Array wrapping the external memory.
    """
    from tack.runtime.kernel_utils import as_address

    shape = _array_shape(shape)
    dtype = as_dtype(dtype)
    if isinstance(dtype, ElementType):
        scalar = dtype.scalar_dtype
        storage = (int(math.prod(shape)) * dtype.size,)
    else:
        scalar, storage = dtype, shape
    backend = _target(device)

    # Validate the pointer's memory space against what this backend expects,
    # on backends that distinguish device memory; Metal MTLBuffer objects
    # and CPU numpy arrays are not addresses and need no check. Any form a
    # device pointer arrives in -- int, numpy integer, CUdeviceptr -- is
    # checked, not only a Python int.
    if backend.device_memory_spaces:
        if as_address(ptr) is None:
            raise TypeError(
                f"array_from_ptr() on the {backend.label} backend takes a "
                f"device address (an integer, or an object int() accepts), "
                f"not {type(ptr).__name__}.")
        space = backend.memory_space(ptr)
        if space not in backend.device_memory_spaces:
            raise ValueError(
                f"Pointer is in '{space}' memory but the device is "
                f"'{backend.label}'. array_from_ptr() requires a device "
                f"pointer. Use tack.array() to copy host data to the device.")

    buf = backend.wrap_ptr(ptr, scalar, storage)
    return _array_view(dtype, shape, buf, writable=writable, device=backend)

tack.memory_space

memory_space(ptr, *, device=None) -> str

Query the memory space of a pointer.

Uses device (by default the default device) to determine where the pointer resides.

Returns:

Type	Description
str	The memory space's name, one of: 'cpu': unregistered host memory (regular malloc/new) 'cuda': CUDA device memory (cudaMalloc) 'cuda_pinned': CUDA pinned host memory (cudaMallocHost) 'cuda_managed': CUDA unified memory (cudaMallocManaged) 'hip': HIP device memory (hipMalloc) 'hip_pinned': HIP pinned host memory (hipHostMalloc) 'hip_managed': HIP unified memory (hipMallocManaged)

Source code in packages/tack-core/src/tack/lang/array.py

def memory_space(ptr, *, device=None) -> str:
    """Query the memory space of a pointer.

    Uses ``device`` (by default the default device) to determine where the
    pointer resides.

    Returns:
        The memory space's name, one of:

        - ``'cpu'``: unregistered host memory (regular malloc/new)
        - ``'cuda'``: CUDA device memory (cudaMalloc)
        - ``'cuda_pinned'``: CUDA pinned host memory (cudaMallocHost)
        - ``'cuda_managed'``: CUDA unified memory (cudaMallocManaged)
        - ``'hip'``: HIP device memory (hipMalloc)
        - ``'hip_pinned'``: HIP pinned host memory (hipHostMalloc)
        - ``'hip_managed'``: HIP unified memory (hipMallocManaged)
    """
    return _target(device).memory_space(ptr)

tack.from_dlpack

from_dlpack(x, *, device=None, copy: bool = False) -> Array

A Tack array from a DLPack capsule or any object with __dlpack__.

Zero-copy: the array shares memory with the source, and the source is held alive for as long as the array lives. Works for device tensors as well as host ones. device must be able to wrap the tensor; by default the array goes on the default device when that can, else on the tensor's own device, as the array API's from_dlpack does.

Wrap, not address: the Metal backend reads host memory perfectly well but cannot turn a host pointer into the MTLBuffer an array needs, so with Metal as the default a host tensor becomes a CPU array, and device="metal" refuses it and says so. Pass copy=True for an independent array instead.

Usage

a = tack.from_dlpack(torch_tensor) a = tack.from_dlpack(cupy_array) a = tack.from_dlpack(numpy_array) a = tack.from_dlpack(vtk_capsule)

Source code in packages/tack-core/src/tack/lang/array.py

def from_dlpack(x, *, device=None, copy: bool = False) -> Array:
    """A Tack array from a DLPack capsule or any object with ``__dlpack__``.

    Zero-copy: the array shares memory with the source, and the source is
    held alive for as long as the array lives. Works for device tensors as
    well as host ones. ``device`` must be able to wrap the tensor; by
    default the array goes on the default device when that can, else on the
    tensor's own device, as the array API's ``from_dlpack`` does.

    Wrap, not address: the Metal backend reads host memory perfectly well
    but cannot turn a host pointer into the MTLBuffer an array needs, so
    with Metal as the default a host tensor becomes a CPU array, and
    ``device="metal"`` refuses it and says so. Pass ``copy=True`` for an
    independent array instead.

    Usage:
        a = tack.from_dlpack(torch_tensor)
        a = tack.from_dlpack(cupy_array)
        a = tack.from_dlpack(numpy_array)
        a = tack.from_dlpack(vtk_capsule)
    """
    if copy:
        host = x if isinstance(x, np.ndarray) else np.from_dlpack(x)
        return array(host, device=device)

    from tack.lang.dlpack import dlpack_to_array
    return dlpack_to_array(x, device=device)

tack.ExportedMemory dataclass

ExportedMemory(
    backend: str,
    size: int,
    allocation_size: int,
    handle: object,
    handle_type: str | None = None,
    device_uuid: bytes | None = None,
)

Handle for sharing GPU memory across APIs (e.g. Tack → Dawn/Vulkan).

Contains only plain Python values — no dependency on any GPU library. The consumer dispatches on handle_type to choose the right import path. backend alone does not settle it: CUDA exports a POSIX fd on Linux and a Win32 KMT handle on Windows, and those are imported by different calls.

tack.lang.dlpack.array_to_dlpack

array_to_dlpack(array, versioned=False)

Create a DLPack capsule from a Tack array.

With versioned, produces a v1.0 DLManagedTensorVersioned in a 'dltensor_versioned' capsule; otherwise the legacy DLManagedTensor in a 'dltensor' one. Only the versioned form can carry the read-only flag, so a non-writable array exported the legacy way arrives as writable -- the legacy protocol simply has no way to say otherwise.

Source code in packages/tack-core/src/tack/lang/dlpack.py

def array_to_dlpack(array, versioned=False):
    """Create a DLPack capsule from a Tack array.

    With `versioned`, produces a v1.0 DLManagedTensorVersioned in a
    'dltensor_versioned' capsule; otherwise the legacy DLManagedTensor in
    a 'dltensor' one. Only the versioned form can carry the read-only
    flag, so a non-writable array exported the legacy way arrives as
    writable -- the legacy protocol simply has no way to say otherwise.
    """
    device_type, device_id, data_ptr = _get_device_info(array)

    dl_dtype = _DTYPE_TO_DLPACK.get(array.scalar_dtype)
    if dl_dtype is None:
        raise TypeError(f"DLPack does not support dtype: {array.dtype}")
    code, bits, lanes = dl_dtype

    shape = array._host_shape
    ndim = len(shape)
    shape_arr = (ctypes.c_int64 * ndim)(*shape)

    # Contiguous C-order strides (in elements, not bytes)
    strides_arr = (ctypes.c_int64 * ndim)()
    stride = 1
    for i in range(ndim - 1, -1, -1):
        strides_arr[i] = stride
        stride *= shape[i]

    if versioned:
        managed = DLManagedTensorVersioned()
        managed.version = DLPackVersion(major=1, minor=0)
        managed.flags = 0 if array._writable else DLPACK_FLAG_BITMASK_READ_ONLY
        capsule_name = b"dltensor_versioned"
        deleter = _dlpack_deleter_versioned
    else:
        managed = DLManagedTensor()
        capsule_name = b"dltensor"
        deleter = _dlpack_deleter

    managed.dl_tensor.data = ctypes.c_void_p(data_ptr)
    managed.dl_tensor.device = DLDevice(device_type=device_type, device_id=device_id)
    managed.dl_tensor.ndim = ndim
    managed.dl_tensor.dtype = DLDataType(code=code, bits=bits, lanes=lanes)
    managed.dl_tensor.shape = shape_arr
    managed.dl_tensor.strides = strides_arr
    managed.dl_tensor.byte_offset = 0

    # The key travels in manager_ctx, which is the only thing the deleter
    # gets back. Pin everything the capsule points at under that key.
    key = next(_ctx_keys)
    managed.manager_ctx = ctypes.c_void_p(key)
    managed.deleter = _natively(deleter)
    _prevent_gc[key] = (array, managed, shape_arr, strides_arr)

    # addressof, not byref: byref yields a temporary whose lifetime is not
    # tied to the capsule. `managed` itself is kept alive by _prevent_gc.
    return _PyCapsule_New(
        ctypes.addressof(managed), capsule_name,
        ctypes.cast(_natively(_capsule_destructor), ctypes.c_void_p),
    )

tack.lang.dlpack.dlpack_to_array

dlpack_to_array(source, writable=True, device=None)

Wrap a DLPack tensor as a Tack array, without copying.

source may be a capsule or anything implementing __dlpack__ -- a CuPy array, a PyTorch tensor, a VTK array via vtkmodules.util.dlpack_support.

The tensor is held for as long as the array lives, so the source may be dropped immediately. The array does not copy, so the two share memory and a write through either is visible to the other.

device is where the array goes. By default it is the default device when that can wrap the tensor, else the tensor's own device, as the array API's from_dlpack keeps an array on its device: CUDA memory imported with the CPU as default is a CUDA array, not a copy.

Raises RuntimeError if device is given and cannot wrap the tensor -- CUDA memory imported onto the CPU device. Metal refuses host tensors for a different reason, and says which: it can read that memory, it just cannot make an MTLBuffer out of it.

Source code in packages/tack-core/src/tack/lang/dlpack.py

def dlpack_to_array(source, writable=True, device=None):
    """Wrap a DLPack tensor as a Tack array, without copying.

    `source` may be a capsule or anything implementing ``__dlpack__`` -- a
    CuPy array, a PyTorch tensor, a VTK array via
    ``vtkmodules.util.dlpack_support``.

    The tensor is held for as long as the array lives, so the source may be
    dropped immediately. The array does not copy, so the two share memory
    and a write through either is visible to the other.

    ``device`` is where the array goes. By default it is the default
    device when that can wrap the tensor, else the tensor's own device, as
    the array API's ``from_dlpack`` keeps an array on its device: CUDA
    memory imported with the CPU as default is a CUDA array, not a copy.

    Raises RuntimeError if ``device`` is given and cannot wrap the tensor --
    CUDA memory imported onto the CPU device. Metal refuses host tensors for
    a different reason, and says which: it can read that memory, it just
    cannot make an MTLBuffer out of it.
    """
    from tack.lang.array import _target, array_from_ptr
    from tack.runtime.dispatch import get_backend

    capsule = _request_capsule(source)

    capsule_ptr = ctypes.cast(id(capsule), ctypes.c_void_p)
    if _PyCapsule_IsValid(capsule_ptr, b"dltensor_versioned"):
        name, used_name = b"dltensor_versioned", b"used_dltensor_versioned"
        struct = DLManagedTensorVersioned
    elif _PyCapsule_IsValid(capsule_ptr, b"dltensor"):
        name, used_name = b"dltensor", b"used_dltensor"
        struct = DLManagedTensor
    else:
        raise ValueError(
            "expected an unconsumed DLPack capsule; this one has already "
            "been taken by another consumer")

    managed_ptr = _PyCapsule_GetPointer(capsule_ptr, name)
    managed = ctypes.cast(managed_ptr, ctypes.POINTER(struct)).contents
    tensor = managed.dl_tensor

    if struct is DLManagedTensorVersioned and \
            managed.flags & DLPACK_FLAG_BITMASK_READ_ONLY:
        # The producer says this memory must not be written. Honour it rather
        # than handing back an array whose writes go somewhere unexpected.
        writable = False

    if not tensor.data:
        raise ValueError("the tensor has a null data pointer")

    dtype = _DLPACK_TO_DTYPE.get((tensor.dtype.code, tensor.dtype.bits))
    if dtype is None:
        raise TypeError(
            f"no Tack dtype for DLPack (code={tensor.dtype.code}, "
            f"bits={tensor.dtype.bits})")
    if tensor.dtype.lanes != 1:
        raise TypeError(
            f"vectorized DLPack dtypes are not supported (lanes={tensor.dtype.lanes})")

    shape = tuple(tensor.shape[i] for i in range(tensor.ndim))

    # Only C-contiguous tensors can be wrapped. Anything else would be read
    # with the wrong stride, which produces plausible-looking wrong numbers.
    if tensor.strides:
        expected, stride = [0] * tensor.ndim, 1
        for i in range(tensor.ndim - 1, -1, -1):
            expected[i] = stride
            stride *= shape[i]
        actual = [tensor.strides[i] for i in range(tensor.ndim)]
        if actual != expected:
            raise ValueError(
                f"only C-contiguous tensors can be wrapped without copying; "
                f"strides {tuple(actual)} != {tuple(expected)}")

    allowed = _DEVICE_BACKENDS.get(tensor.device.device_type)
    if tensor.device.device_type == kDLMetal:
        raise ValueError(
            "Metal DLPack tensors (kDLMetal) cannot be imported: Tack has "
            "no way yet to wrap the MTLBuffer handle they carry. Move the "
            "tensor to host memory in the producing library, then import "
            "it with copy=True.")
    if allowed is None:
        raise ValueError(
            f"unsupported DLPack device type {tensor.device.device_type}")
    where = (tensor.device.device_type, tensor.device.device_id)
    if device is not None:
        backend = _target(device)
    else:
        backend = get_backend()
        if not backend.addresses_dlpack(*where):
            backend = _tensor_device(allowed[0], tensor.device.device_id)
    if backend.name not in allowed and not backend.addresses_dlpack(*where):
        message = (
            f"the tensor is on a device the '{backend.name}' device cannot "
            f"wrap without copying (DLPack device type "
            f"{tensor.device.device_type}). Import it on its own device with "
            f"tack.from_dlpack(x, device='{allowed[0]}'), or pass copy=True.")
        if backend.dlpack_refusal_note:
            message += f" {backend.dlpack_refusal_note}"
        raise RuntimeError(message)

    array = array_from_ptr(
        (tensor.data or 0) + tensor.byte_offset, shape, dtype, device=backend,
        writable=writable)

    # Adopt the capsule: rename it so nobody else releases it, and attach the
    # hold to the array so the producer is told when the array is gone.
    _PyCapsule_SetName(capsule, used_name)
    array._dlpack_hold = _CapsuleHold(capsule, managed_ptr, struct)
    return array

tack.lang.dlpack.dlpack_device

dlpack_device(array)

Return (device_type, device_id) tuple for DLPack protocol.

Source code in packages/tack-core/src/tack/lang/dlpack.py

def dlpack_device(array):
    """Return (device_type, device_id) tuple for DLPack protocol."""
    device_type, device_id, _ = _get_device_info(array)
    return (device_type, device_id)

Kernels and device functions

tack.kernel

kernel(func)

Decorator that marks a Python function as a GPU kernel.

Source code in packages/tack-core/src/tack/lang/kernel.py

def kernel(func):
    """Decorator that marks a Python function as a GPU kernel."""
    return Kernel(func)

tack.lang.kernel.Kernel

Kernel(func, funcdef=None)

A captured kernel function, ready for AST transformation and compilation.

Source code in packages/tack-core/src/tack/lang/kernel.py

def __init__(self, func, funcdef=None):
    self.func = func
    if funcdef is None:
        self._source = textwrap.dedent(read_source(func, "kernel"))
        self._ast = ast.parse(self._source)
        tag_origin(self._ast, func)
        self._funcdef = self._ast.body[0]  # The FunctionDef node
    else:
        # A kernel derived from another's AST: the result epilogue below.
        self._source = ast.unparse(funcdef)
        self._ast = ast.Module(body=[funcdef], type_ignores=[])
        self._funcdef = funcdef
    self.name = self._funcdef.name
    # Lazy IR: defer transform until first dispatch (vector arrays may be needed)
    self._ir = None
    self._ir_cache = {}  # vector_arrays key → IRModule
    # A kernel ending in `return expr` has the expression computed by a
    # one-thread epilogue kernel into a hidden array; see _split_result.
    self._result = None if funcdef is not None else _split_result(self, func)
    # Per thread: each call stores its result there and reads it back,
    # so two threads sharing one array read each other's result.
    self._result_arrays = threading.local()

get_ir

get_ir(
    vector_arrays=None,
    template_args=None,
    texture_arrays=None,
    debug=False,
)

Get IR, re-transforming if vector/texture array or template metadata is provided.

debug lowers with bounds checks (see ir_bounds), cached apart.

Source code in packages/tack-core/src/tack/lang/kernel.py

def get_ir(self, vector_arrays=None, template_args=None, texture_arrays=None, debug=False):
    """Get IR, re-transforming if vector/texture array or template metadata is provided.

    ``debug`` lowers with bounds checks (see ``ir_bounds``), cached apart.
    """
    if template_args:
        key = self._make_cache_key(vector_arrays, template_args, texture_arrays, debug)
        cached = self._ir_cache.get(key)
        if cached is not None:
            return cached
        with _transform_lock:
            # Re-check: another thread may have built it while we waited.
            if key not in self._ir_cache:
                self._ir_cache[key] = self._lower(vector_arrays, texture_arrays,
                                                  template_args, debug)
                for token in _key_tokens(key):
                    token.kernels.add(self)
            return self._ir_cache[key]
    if not vector_arrays and not texture_arrays and not debug:
        if self._ir is None:
            with _transform_lock:
                if self._ir is None:
                    self._ir = self._lower(None, None, None)
        return self._ir
    key = self._make_cache_key(vector_arrays, None, texture_arrays, debug)
    cached = self._ir_cache.get(key)
    if cached is not None:
        return cached
    with _transform_lock:
        if key not in self._ir_cache:
            self._ir_cache[key] = self._lower(vector_arrays, texture_arrays, None, debug)
        return self._ir_cache[key]

__get__

__get__(instance, owner=None)

Bind a kernel defined in a class body to the instance it is called on.

model.step(dt) passes model as the kernel's first argument, where a @tack.device_class object is a template like any other. Reached through the class, the kernel is unbound.

Source code in packages/tack-core/src/tack/lang/kernel.py

def __get__(self, instance, owner=None):
    """Bind a kernel defined in a class body to the instance it is called on.

    ``model.step(dt)`` passes ``model`` as the kernel's first argument,
    where a ``@tack.device_class`` object is a template like any
    other. Reached through the class, the kernel is unbound.
    """
    if instance is None:
        return self
    return types.MethodType(self, instance)

tack.ndindex

ndindex(*args)

Multi-dimensional parallel iteration, as np.ndindex.

Used in kernels: for i, j in tack.ndindex(w, h); each argument is a size or a (start, end) pair. Cannot be called from Python directly.

Source code in packages/tack-core/src/tack/__init__.py

def ndindex(*args):
    """Multi-dimensional parallel iteration, as ``np.ndindex``.

    Used in kernels: ``for i, j in tack.ndindex(w, h)``; each argument is a
    size or a ``(start, end)`` pair. Cannot be called from Python directly.
    """
    raise RuntimeError("ndindex() can only be used inside a @tack.kernel")

tack.func

func(f)

Decorator that marks a Python function as a device-side inlineable function.

Source code in packages/tack-core/src/tack/lang/func.py

def func(f):
    """Decorator that marks a Python function as a device-side inlineable function."""
    return Func(f)

tack.lang.func.Func

Func(func)

A captured device-side function, ready for AST inlining.

Source code in packages/tack-core/src/tack/lang/func.py

def __init__(self, func):
    self.func = func
    self.name = func.__name__
    self._source = textwrap.dedent(read_source(func, "device function"))
    self._ast = ast.parse(self._source)
    tag_origin(self._ast, func)
    # Extract the FunctionDef node
    self._funcdef = self._ast.body[0]
    if not isinstance(self._funcdef, ast.FunctionDef):
        raise TypeError(f"@tack.func must decorate a function, got {type(self._funcdef)}")
    # Class methods are collected by @tack.device_class and resolved
    # per template transformation. Ordinary functions need no registry.
    self._is_method = (
        len(self._funcdef.args.args) > 0 and
        self._funcdef.args.args[0].arg == 'self'
    )

tack.device_class

device_class(cls)

Mark a class as a device class: its objects can be kernel arguments.

Methods decorated with @tack.func are collected and made available for inlining when instances are passed to kernels.

Source code in packages/tack-core/src/tack/lang/device_class.py

def device_class(cls):
    """Mark a class as a device class: its objects can be kernel arguments.

    Methods decorated with @tack.func are collected and made available
    for inlining when instances are passed to kernels.
    """
    from tack.lang.template_rewrite import template_func_methods

    cls._device_class = True
    # A record of the methods at decoration, inherited ones included.
    # Lowering looks them up again per class, so subclasses defined later,
    # decorated or not, get their own.
    cls._tack_func_methods = template_func_methods(cls)

    return cls

tack.constant

constant(value, dtype=None)

Declare a named constant that kernels and device functions may read.

DT = tack.constant(0.01) at module (or enclosing function) scope lets a kernel write DT as if it were the literal 0.01. Without dtype the constant behaves exactly like that literal: an integer takes the narrowest of i32, i64 and u64 that holds it, and a float is weak, taking the precision of what it meets. With dtype it has that type, as tack.u32(747796405) or tack.f64(0.1) written in the kernel would: tack.constant(747796405, tack.u32) multiplies in wrapping u32 arithmetic, and tack.constant(0.1, tack.f64) is the exact double.

The value is fixed here. A kernel reads the constant its name is bound to when the kernel is first compiled; binding the name to another constant afterwards does not change a compiled kernel, the same rule as for device functions.

Source code in packages/tack-core/src/tack/lang/constant.py

def constant(value, dtype=None):
    """Declare a named constant that kernels and device functions may read.

    ``DT = tack.constant(0.01)`` at module (or enclosing function) scope
    lets a kernel write ``DT`` as if it were the literal ``0.01``. Without
    ``dtype`` the constant behaves exactly like that literal: an integer
    takes the narrowest of i32, i64 and u64 that holds it, and a float is
    weak, taking the precision of what it meets. With ``dtype`` it has that
    type, as ``tack.u32(747796405)`` or ``tack.f64(0.1)`` written in the
    kernel would: ``tack.constant(747796405, tack.u32)`` multiplies in
    wrapping u32 arithmetic, and ``tack.constant(0.1, tack.f64)`` is the
    exact double.

    The value is fixed here. A kernel reads the constant its name is bound
    to when the kernel is first compiled; binding the name to another
    constant afterwards does not change a compiled kernel, the same rule
    as for device functions.
    """
    if dtype is not None and not isinstance(dtype, ScalarType):
        raise TypeError(f"tack.constant dtype must be a Tack scalar type, not {dtype!r}")
    if isinstance(value, (tuple, list, np.ndarray)):
        return _array_constant(value, dtype)
    if isinstance(value, (bool, numbers.Integral)):
        value = int(value)
        if dtype is None or dtype in INTEGER_TYPES:
            _check_representable(value, dtype)
            return IntConstant(value, dtype)
        return FloatConstant(float(value), dtype)
    if isinstance(value, numbers.Real):
        if dtype in INTEGER_TYPES:
            raise TypeError(
                f"tack.constant({value!r}, {dtype!r}): an integer type needs an integer "
                f"value; convert it explicitly with int() if truncation is intended")
        return FloatConstant(float(value), dtype)
    raise TypeError(
        f"tack.constant takes an int, a float, or a tuple of them (a vector) or of "
        f"tuples (a matrix), not {type(value).__name__}; arrays and other objects are "
        f"passed to kernels as arguments")

tack.lang.constant.ArrayConstant

ArrayConstant(components, shape, dtype)

A vector or matrix that kernels may read by name.

SUN = tack.constant((0.5, 0.5, 0.0)) is the vector tack.vec3(0.5, 0.5, 0.0) in a kernel, and tack.constant(((1, 0), (0, 1))) the matrix with those rows, each component a constant with the rules of tack.constant. On the host it indexes and iterates like a tuple of its rows and converts to a NumPy array: np.array(SUN, np.float32), SUN[1], len(SUN). It is not a tuple, so + does not concatenate it by mistake; do host arithmetic on np.asarray(SUN).

Source code in packages/tack-core/src/tack/lang/constant.py

def __init__(self, components, shape, dtype):
    self.components = tuple(components)      # flat, row-major
    self.shape = tuple(shape)
    self.dtype = dtype

Shader-style helpers

tack.math

Shader-style helpers as device functions: fract, mix, clamp...

GLSL, MSL and Taichi's taichi.math share a small vocabulary that graphics and simulation code reaches for constantly, and that Python's math lacks. These are ordinary @tack.func device functions, so they inline into kernels, and because Tack's arithmetic and builtins apply to each component of a vector, each one works on scalars and on vectors (and matrices) alike, as the GLSL functions do::

from tack import math as tm

@tack.kernel
def shade(uv, out, n):
    for i in range(n):
        st = tm.fract(uv[i] * 8.0) - 0.5          # a 2-vector
        edge = tm.smoothstep(0.0, 0.05, st.x)     # a scalar
        out[i] = tm.mix(DARK, LIGHT, edge)        # 3-vectors

The definitions are GLSL's, so mix(x, y, a) is x * (1 - a) + y * a, smoothstep is the Hermite polynomial on the clamped ratio, and step and sign select per component through tack.select. They are only callable from kernels and device functions; a NumPy reference writes the same one-liners with np.floor, np.clip and friends.

fract

fract(x)

The fractional part, x - floor(x), in [0, 1).

Source code in packages/tack-core/src/tack/math.py

@tack.func
def fract(x):
    """The fractional part, ``x - floor(x)``, in ``[0, 1)``."""
    return x - floor(x)

mix

mix(x, y, a)

The linear blend x * (1 - a) + y * a; a may be a scalar or match x.

Source code in packages/tack-core/src/tack/math.py

@tack.func
def mix(x, y, a):
    """The linear blend ``x * (1 - a) + y * a``; ``a`` may be a scalar or match ``x``."""
    return x * (1.0 - a) + y * a

clamp

clamp(x, lo, hi)

x limited to [lo, hi], component by component.

Source code in packages/tack-core/src/tack/math.py

@tack.func
def clamp(x, lo, hi):
    """``x`` limited to ``[lo, hi]``, component by component."""
    return min(max(x, lo), hi)

saturate

saturate(x)

x limited to [0, 1].

Source code in packages/tack-core/src/tack/math.py

@tack.func
def saturate(x):
    """``x`` limited to ``[0, 1]``."""
    return min(max(x, 0.0), 1.0)

smoothstep

smoothstep(edge0, edge1, x)

Hermite interpolation: 0 below edge0, 1 above edge1, smooth between.

edge0 greater than edge1 reverses the step, as in GLSL.

Source code in packages/tack-core/src/tack/math.py

@tack.func
def smoothstep(edge0, edge1, x):
    """Hermite interpolation: 0 below ``edge0``, 1 above ``edge1``, smooth between.

    ``edge0`` greater than ``edge1`` reverses the step, as in GLSL.
    """
    t = min(max((x - edge0) / (edge1 - edge0), 0.0), 1.0)
    return t * t * (3.0 - 2.0 * t)

step

step(edge, x)

0.0 where x < edge, else 1.0, component by component.

Source code in packages/tack-core/src/tack/math.py

@tack.func
def step(edge, x):
    """0.0 where ``x < edge``, else 1.0, component by component."""
    return tack.select(x < edge, 0.0, 1.0)

sign

sign(x)

-1.0, 0.0 or 1.0 by the sign of x, component by component.

Source code in packages/tack-core/src/tack/math.py

@tack.func
def sign(x):
    """-1.0, 0.0 or 1.0 by the sign of ``x``, component by component."""
    return tack.select(x > 0.0, 1.0, tack.select(x < 0.0, -1.0, 0.0))

length

length(v)

The Euclidean length of a vector.

Source code in packages/tack-core/src/tack/math.py

@tack.func
def length(v):
    """The Euclidean length of a vector."""
    return v.norm()

distance

distance(a, b)

The Euclidean distance between two points.

Source code in packages/tack-core/src/tack/math.py

@tack.func
def distance(a, b):
    """The Euclidean distance between two points."""
    return (a - b).norm()

normalize

normalize(v)

v scaled to unit length; a zero vector gives NaNs, as in GLSL.

Source code in packages/tack-core/src/tack/math.py

@tack.func
def normalize(v):
    """``v`` scaled to unit length; a zero vector gives NaNs, as in GLSL."""
    return v / v.norm()

Random numbers

tack.random

Random numbers in kernels, from a counter-based generator with explicit state.

Nothing here is hidden: a kernel seeds a u32 state from the element it works on and a stream number, and every draw returns the value and the next state. The sequence is a pure function of those integers, so the same program draws the same numbers on every backend and for any thread order, and a NumPy reference that calls the np_ mirrors below draws exactly the same ones::

from tack import random

@tack.kernel
def scatter(pos, n, frame):
    for i in range(n):
        state = random.seed(i, frame)        # element i of this frame
        u, state = random.uniform(state)     # f32 in [0, 1)
        z, state = random.normal(state)      # standard normal
        d, state = random.direction3(state)  # a unit 3-vector
        pos[i] = [u, z, d.x]

The generator is PCG's output permutation (RXS-M-XS on 32 bits) over a linear congruential step, in wrapping u32 arithmetic. uniform is the state's top 24 bits scaled, so it is exact and identical everywhere; normal and the directions go through log, sqrt, cos and sin, so they agree across backends to those functions' rounding. Taichi's ti.random() keeps hidden per-thread state instead; its results depend on the thread schedule and cannot be reproduced by a reference.

advance

advance(state)

The state after state; also the 32 random bits it stands for.

Source code in packages/tack-core/src/tack/random.py

@tack.func
def advance(state):
    """The state after ``state``; also the 32 random bits it stands for."""
    x = tack.u32(state) * MULTIPLIER + INCREMENT
    word = ((x >> ((x >> S28) + S4)) ^ x) * MIX
    return (word >> S22) ^ word

seed

seed(index, stream)

A state for element index of stream (a frame, a pass, a purpose).

Source code in packages/tack-core/src/tack/random.py

@tack.func
def seed(index, stream):
    """A state for element ``index`` of ``stream`` (a frame, a pass, a purpose)."""
    return advance(tack.u32(index) ^ advance(tack.u32(stream)))

uniform

uniform(state)

An f32 in [0, 1) and the next state.

Source code in packages/tack-core/src/tack/random.py

@tack.func
def uniform(state):
    """An f32 in [0, 1) and the next state."""
    return tack.f32(state >> S8) * SCALE, advance(state)

normal

normal(state)

A standard normal draw (Box-Muller) and the next state; uses two states.

Source code in packages/tack-core/src/tack/random.py

@tack.func
def normal(state):
    """A standard normal draw (Box-Muller) and the next state; uses two states."""
    u1 = 1.0 - tack.f32(state >> S8) * SCALE              # in (0, 1], so log is finite
    state = advance(state)
    u2 = tack.f32(state >> S8) * SCALE
    return sqrt(-2.0 * log(u1)) * cos(TWO_PI * u2), advance(state)

direction2

direction2(state)

A unit 2-vector with uniformly distributed direction, and the next state.

Source code in packages/tack-core/src/tack/random.py

@tack.func
def direction2(state):
    """A unit 2-vector with uniformly distributed direction, and the next state."""
    a = TWO_PI * tack.f32(state >> S8) * SCALE
    return tack.vec2(cos(a), sin(a)), advance(state)

direction3

direction3(state)

A unit 3-vector uniformly distributed on the sphere, and the next state; two states.

Source code in packages/tack-core/src/tack/random.py

@tack.func
def direction3(state):
    """A unit 3-vector uniformly distributed on the sphere, and the next state; two states."""
    a = TWO_PI * tack.f32(state >> S8) * SCALE
    state = advance(state)
    z = tack.f32(state >> S8) * SCALE * 2.0 - 1.0
    r = sqrt(1.0 - z * z)
    return tack.vec3(r * cos(a), r * sin(a), z), advance(state)

tack.lang.source_validation.UnsupportedSyntaxError

Bases: NotImplementedError

A source construct outside Tack's kernel language.

Inspection

tack.inspect

inspect(kernel, *args, mode='source')

Return generated code for a kernel without executing it.

Parameters:

Name	Type	Description	Default
kernel		A @tack.kernel decorated function.	required
*args		The arguments the kernel would be called with.	()
mode		What to return: "ir" — Tack intermediate representation "source" — backend-specific source code (MSL, CUDA C, LLVM IR, etc.) "optimized" — LLVM IR after O3; CPU only, since the GPU backends' vendor compilers optimize past anything Tack sees	'source'

Returns:

Type	Description
	The generated code as a string.

Source code in packages/tack-core/src/tack/lang/inspect_kernel.py

def inspect(kernel, *args, mode="source"):
    """Return generated code for a kernel without executing it.

    Args:
        kernel: A @tack.kernel decorated function.
        *args: The arguments the kernel would be called with.
        mode: What to return:
            "ir"        — Tack intermediate representation
            "source"    — backend-specific source code (MSL, CUDA C, LLVM IR, etc.)
            "optimized" — LLVM IR after O3; CPU only, since the GPU backends'
                          vendor compilers optimize past anything Tack sees

    Returns:
        The generated code as a string.
    """
    from tack.lang.kernel import Kernel
    bound_to = getattr(kernel, '__self__', None)
    if bound_to is not None and isinstance(getattr(kernel, '__func__', None), Kernel):
        # A kernel method reached through its instance: model.step
        kernel, args = kernel.__func__, (bound_to, *args)
    if not isinstance(kernel, Kernel):
        raise TypeError(f"Expected a @tack.kernel, got {type(kernel).__name__}")

    from tack.runtime.dispatch import launch_backend
    backend = launch_backend(kernel.name, args)
    if not backend.supports_inspect:
        raise NotImplementedError(
            f"tack.inspect does not show kernels on the {backend.label} backend; "
            f"inspect '{kernel.name}' with arrays on another device")

    if mode == "ir":
        ir_func, _ = _prepare_ir(kernel, args, backend=backend)
        return ir.dump(ir_func)

    if mode == "source":
        return _generate_source(kernel, args)

    if mode == "optimized":
        return _generate_source(kernel, args, optimize=True)

    raise ValueError(
        f"Unknown inspect mode: '{mode}'. Use 'ir', 'source', or 'optimized'."
    )

Reductions, statistics and scans

How these behave on each backend — dtypes, accumulators, NaN and empty inputs, repeatability — is in the User's Guide chapter Reductions and Scans. Array reductions are the methods Array.sum, min, max and mean, above.

Statistics

Import these from tack.algorithms; they are defined in tack.algorithms.stats.

tack.algorithms.stats.var

var(data, n=None)

Population variance of an array: Σ(x - mean)² / n.

Two passes over the first n elements: the mean by Array.sum(), then the squared differences, formed in the data's precision (f64 for f64 data, otherwise f32) and summed by Array.sum(). Returns a Python float; an empty range gives NaN, as Array.mean() does.

Source code in packages/tack-core/src/tack/algorithms/stats.py

@on_inputs_device
def var(data, n=None):
    """Population variance of an array: Σ(x - mean)² / n.

    Two passes over the first n elements: the mean by ``Array.sum()``, then
    the squared differences, formed in the data's precision (f64 for f64
    data, otherwise f32) and summed by ``Array.sum()``. Returns a Python
    float; an empty range gives NaN, as ``Array.mean()`` does.
    """
    n, data = _count(n, data)
    if n == 0:
        return float('nan')
    if data.dtype in (tack.f32, tack.f64):
        mean_val = _prefix(data, n).sum() / n
    else:
        # An integer sum wraps (2^62 four times is 0): the values convert
        # first, as every other term does.
        mean_val = _terms(_converted, n, data).sum() / n
    terms = tack.empty((n,), _precision(data))
    _squared_diffs(data, mean_val, 1.0, terms, n)
    return terms.sum() / n

tack.algorithms.stats.std

std(data, n=None)

Population standard deviation of an array.

Source code in packages/tack-core/src/tack/algorithms/stats.py

@on_inputs_device
def std(data, n=None):
    """Population standard deviation of an array."""
    from math import sqrt
    return sqrt(var(data, n))

tack.algorithms.stats.norm

norm(data, ord=2, n=None)

Vector norm of an array, in the data's precision (f64 for f64 data, otherwise f32), as a Python float.

ord=1: L1 norm (sum of absolute values) ord=2: L2 norm (Euclidean) ord=inf: L-infinity (max absolute value)

Any other ord raises ValueError. An empty range gives 0.0.

Source code in packages/tack-core/src/tack/algorithms/stats.py

@on_inputs_device
def norm(data, ord=2, n=None):
    """Vector norm of an array, in the data's precision (f64 for f64 data,
    otherwise f32), as a Python float.

    ord=1: L1 norm (sum of absolute values)
    ord=2: L2 norm (Euclidean)
    ord=inf: L-infinity (max absolute value)

    Any other ord raises ValueError. An empty range gives 0.0.
    """
    if ord not in (1, 2, float('inf')):
        raise ValueError(f"Unsupported norm order: {ord}")
    n, data = _count(n, data)
    if n == 0:
        return 0.0
    if ord == 1:
        return _terms(_abs_values, n, data).sum()
    if ord == 2:
        from math import sqrt
        return sqrt(_terms(_squares, n, data).sum())
    return _terms(_abs_values, n, data).max()

tack.algorithms.stats.absmax

absmax(data, n=None)

Maximum absolute value of an array, as a Python float; 0.0 when empty.

Source code in packages/tack-core/src/tack/algorithms/stats.py

@on_inputs_device
def absmax(data, n=None):
    """Maximum absolute value of an array, as a Python float; 0.0 when empty."""
    return norm(data, float('inf'), n)

tack.algorithms.stats.count_nonzero

count_nonzero(data, n=None)

The number of nonzero elements of an array, as a Python int: its 0/1 flags scanned into i64 on the device, so the count is exact past 2^31 and only the total comes back to the host.

Source code in packages/tack-core/src/tack/algorithms/stats.py

@on_inputs_device
def count_nonzero(data, n=None):
    """The number of nonzero elements of an array, as a Python int: its 0/1 flags
    scanned into i64 on the device, so the count is exact past 2^31 and only
    the total comes back to the host."""
    from tack.algorithms.scan import exclusive_scan

    n, data = _count(n, data)
    if n == 0:
        return 0
    flags = tack.empty((n,), tack.i32)
    _nonzero(data, flags, n)
    return int(exclusive_scan(flags, tack.empty((n,), tack.i64), n))

tack.algorithms.stats.dot

dot(a, b, n=None)

Dot product of two arrays, Σa[i]·b[i], in their precision (f64 if either is f64, otherwise f32), as a Python float; 0.0 when empty.

Source code in packages/tack-core/src/tack/algorithms/stats.py

@on_inputs_device
def dot(a, b, n=None):
    """Dot product of two arrays, Σa[i]·b[i], in their precision (f64 if either
    is f64, otherwise f32), as a Python float; 0.0 when empty."""
    n, a, b = _count(n, a, b)
    if n == 0:
        return 0.0
    return _terms(_products, n, a, b).sum()

tack.algorithms.stats.histogram

histogram(data, bins=10, range=None, n=None)

Compute a histogram of array values with atomic bin counts.

Values below the range are counted in the first bin and values above it in the last, unlike numpy.histogram, which drops them.

Parameters:

Name	Type	Description	Default
data		input array of any dtype	required
bins		number of bins (at least 1)	10
range		(min, max) tuple. If None, uses the minimum and maximum of the first n elements, and n must then be at least 1.	None
n		number of elements (default: data.size)	None

Returns:

Type	Description
	(counts, bin_edges) where counts is a Tack array of i32,
	bin_edges is a numpy array of (bins + 1) float64 edges.

Source code in packages/tack-core/src/tack/algorithms/stats.py

@on_inputs_device
def histogram(data, bins=10, range=None, n=None):
    """Compute a histogram of array values with atomic bin counts.

    Values below the range are counted in the first bin and values above it
    in the last, unlike numpy.histogram, which drops them.

    Args:
        data: input array of any dtype
        bins: number of bins (at least 1)
        range: (min, max) tuple. If None, uses the minimum and maximum of
            the first n elements, and n must then be at least 1.
        n: number of elements (default: data.size)

    Returns:
        (counts, bin_edges) where counts is a Tack array of i32,
        bin_edges is a numpy array of (bins + 1) float64 edges.
    """
    import numpy as np
    n, data = _count(n, data)
    if int(bins) < 1:
        raise ValueError(f"histogram needs at least 1 bin, not {bins}")
    if range is None:
        if n == 0:
            raise ValueError("histogram of no elements needs an explicit range")
        head = _prefix(data, n)
        lo = float(head.min())
        hi = float(head.max())
        if not (np.isfinite(lo) and np.isfinite(hi)):
            # As numpy.histogram: bins cannot divide a range that is not finite.
            raise ValueError(f"autodetected range of [{lo}, {hi}] is not finite; "
                             "pass a finite range")
    else:
        lo, hi = float(range[0]), float(range[1])
        if not (np.isfinite(lo) and np.isfinite(hi)):
            raise ValueError(f"histogram range [{lo}, {hi}] is not finite")
        if lo > hi:
            raise ValueError(f"histogram range [{lo}, {hi}] has its minimum above its maximum")

    # Avoid division by zero for constant arrays
    if hi == lo:
        hi = lo + 1.0

    bin_width = (hi - lo) / bins
    inv_bw = 1.0 / bin_width

    counts = tack.empty((bins,), tack.i32)
    counts.fill(0)
    _histogram_kernel(data, counts, lo, inv_bw, bins, n)

    edges = np.linspace(lo, hi, bins + 1)
    return counts, edges

Scans

Import these from tack.algorithms; they are defined in tack.algorithms.scan.

tack.algorithms.scan.exclusive_scan

exclusive_scan(
    input_array, output_array, n, total=True, close=False
)

Compute exclusive prefix sum on the inputs' device.

output[i] = sum(input[0..i-1]), output[0] = 0.

The scan runs in the output array's dtype; input values convert to it on the copy, so any dtype the backend allocates works.

Parameters:

Name	Type	Description	Default
input_array		array with input values.	required
output_array		array for the output offsets.	required
n		number of elements, at most either array's size. Zero writes nothing and returns 0.	required
total		whether to return the total. Reading it waits for the device (on Metal, for every launch queued so far); a caller that needs only the offsets passes total=False and gets None, and the scan stays queued.	True
close		also write the total at output[n], so output is the n + 1 offsets of CSR rows (output[0] is 0 when n is 0). In the scan's own launch where it is one, as for any array of at most a tile or 4096 elements; output needs n + 1 elements.	False

Returns:

Type	Description
	The total of all input elements in the output's dtype, as a Python
	int or float; None with total=False.

Source code in packages/tack-core/src/tack/algorithms/scan.py

@on_inputs_device
def exclusive_scan(input_array, output_array, n, total=True, close=False):
    """Compute exclusive prefix sum on the inputs' device.

    output[i] = sum(input[0..i-1]), output[0] = 0.

    The scan runs in the output array's dtype; input values convert to it
    on the copy, so any dtype the backend allocates works.

    Args:
        input_array: array with input values.
        output_array: array for the output offsets.
        n: number of elements, at most either array's size. Zero writes
            nothing and returns 0.
        total: whether to return the total. Reading it waits for the
            device (on Metal, for every launch queued so far); a caller
            that needs only the offsets passes ``total=False`` and gets
            ``None``, and the scan stays queued.
        close: also write the total at ``output[n]``, so ``output`` is the
            ``n + 1`` offsets of CSR rows (``output[0]`` is 0 when ``n`` is
            0). In the scan's own launch where it is one, as for any array
            of at most a tile or 4096 elements; ``output`` needs ``n + 1``
            elements.

    Returns:
        The total of all input elements in the output's dtype, as a Python
        int or float; ``None`` with ``total=False``.
    """
    n = _count(n, input_array, output_array)
    if close and output_array.size < n + 1:
        raise ValueError(f"exclusive_scan with close writes output[{n}]: the output has "
                         f"{output_array.size} elements")
    if n == 0:
        # Nothing to scan; the empty sum is 0, as for Array.sum().
        if close:
            _close_empty(output_array)
        return 0 if total else None
    # Scan in the output's dtype, as the inclusive scan does: an i32 work
    # buffer truncated float inputs and wrapped wider integers.
    return _scan(input_array, output_array, n, inclusive=False, want_total=total, close=close)

tack.algorithms.scan.inclusive_scan

inclusive_scan(input_array, output_array, n, total=True)

Compute inclusive prefix sum on the inputs' device.

output[i] = sum(input[0..i]).

The scan runs in the output array's dtype; input values convert to it on the copy, so any dtype the backend allocates works.

Parameters:

Name	Type	Description	Default
input_array		array with input values.	required
output_array		array for the output sums.	required
n		number of elements, at most either array's size. Zero writes nothing and returns 0.	required
total		whether to return the total; see exclusive_scan.	True

Returns:

Type	Description
	The total of all input elements in the output's dtype, as a Python
	int or float; None with total=False.

Source code in packages/tack-core/src/tack/algorithms/scan.py

@on_inputs_device
def inclusive_scan(input_array, output_array, n, total=True):
    """Compute inclusive prefix sum on the inputs' device.

    output[i] = sum(input[0..i]).

    The scan runs in the output array's dtype; input values convert to it
    on the copy, so any dtype the backend allocates works.

    Args:
        input_array: array with input values.
        output_array: array for the output sums.
        n: number of elements, at most either array's size. Zero writes
            nothing and returns 0.
        total: whether to return the total; see ``exclusive_scan``.

    Returns:
        The total of all input elements in the output's dtype, as a Python
        int or float; ``None`` with ``total=False``.
    """
    n = _count(n, input_array, output_array)
    if n == 0:
        return 0 if total else None
    return _scan(input_array, output_array, n, inclusive=True, want_total=total)

Sorting and segmented reductions

Import these from tack.algorithms; they are defined in tack.algorithms.sort.

tack.algorithms.sort.argsort

argsort(keys, n=None, index_dtype=None)

Stable permutation that sorts the first n keys ascending.

keys is an i32, u32, i64 or u64 array. Returns a new array perm of n elements with keys[perm[0]] <= keys[perm[1]] <= ...; equal keys keep their original order. n defaults to the whole array and may be zero. perm is i32, or i64 past 2**31 - 1 elements or when index_dtype asks.

Source code in packages/tack-core/src/tack/algorithms/sort.py

@on_inputs_device
def argsort(keys, n=None, index_dtype=None):
    """Stable permutation that sorts the first ``n`` keys ascending.

    ``keys`` is an ``i32``, ``u32``, ``i64`` or ``u64`` array.  Returns a
    new array ``perm`` of ``n`` elements with ``keys[perm[0]] <=
    keys[perm[1]] <= ...``; equal keys keep their original order.  ``n``
    defaults to the whole array and may be zero.  ``perm`` is ``i32``, or
    ``i64`` past ``2**31 - 1`` elements or when ``index_dtype`` asks.
    """
    _key_width(keys)
    n = _count(n, keys)
    indices = _index_dtype(n, index_dtype)
    perm = tack.empty((n,), indices)
    if n == 0:
        return perm
    _iota(perm, n)
    if n == 1:
        return perm

    flip = 0 if keys.dtype in (tack.u32, tack.u64) else 1 << (keys.dtype.bits - 1)
    mapped = tack.empty((n,), tack.u64)
    map_keys = _map_keys_32 if keys.dtype.bits == 32 else _map_keys_64
    map_keys(keys, mapped, flip, n)

    # The passes needed depend on the spread of the keys, not their width.
    # The base is the smallest mapped key, so digits above the spread are
    # all zero and those passes can be skipped.
    base, largest = _key_range(mapped, n)
    spread = largest - base
    passes = (spread.bit_length() + _RADIX_BITS - 1) // _RADIX_BITS
    if passes == 0:
        return perm  # every key is equal: the identity is the stable order

    from tack.runtime.dispatch import get_backend

    mapped_b = tack.empty((n,), tack.u64)
    perm_b = tack.empty((n,), indices)
    if get_backend().supports_workgroups:
        # Tiles sorted into the second buffers, then scattered back: the
        # pass's result is where its input was.
        ntiles = (n + _TILE - 1) // _TILE
        counts = tack.empty((_RADIX * ntiles,), tack.i32)
        offsets = tack.empty((_RADIX * ntiles,), indices)
        starts = tack.empty((_RADIX * ntiles,), tack.i32)
        for p in range(passes):
            shift = p * _RADIX_BITS
            _sort_tiles(mapped, perm, mapped_b, perm_b, counts, starts, base, shift, n,
                        ntiles)
            exclusive_scan(counts, offsets, _RADIX * ntiles, total=False)
            _scatter_tiles(mapped_b, perm_b, mapped, perm, offsets, starts, base, shift, n,
                           ntiles)
        return perm
    nchunks = (n + _CHUNK - 1) // _CHUNK
    counts = tack.empty((_RADIX * nchunks,), tack.i32)
    offsets = tack.empty((_RADIX * nchunks,), indices)
    src_keys, src_perm, dst_keys, dst_perm = mapped, perm, mapped_b, perm_b
    for p in range(passes):
        shift = p * _RADIX_BITS
        _count_digits(src_keys, counts, base, shift, n, _CHUNK, nchunks)
        exclusive_scan(counts, offsets, _RADIX * nchunks, total=False)
        _scatter_digits(src_keys, src_perm, dst_keys, dst_perm, offsets,
                        base, shift, n, _CHUNK, nchunks)
        src_keys, dst_keys = dst_keys, src_keys
        src_perm, dst_perm = dst_perm, src_perm
    if src_perm is not perm:
        copy(src_perm, perm, n)
    return perm

tack.algorithms.sort.sort_by_key

sort_by_key(keys, values=None, n=None)

Sort the first n keys ascending, carrying values along.

Returns (sorted_keys, sorted_values) as new arrays of n elements; sorted_values is None when no values are given. The sort is stable, so equal keys keep the order of their values. See argsort for the key dtypes.

Source code in packages/tack-core/src/tack/algorithms/sort.py

@on_inputs_device
def sort_by_key(keys, values=None, n=None):
    """Sort the first ``n`` keys ascending, carrying ``values`` along.

    Returns ``(sorted_keys, sorted_values)`` as new arrays of ``n``
    elements; ``sorted_values`` is ``None`` when no values are given.  The
    sort is stable, so equal keys keep the order of their values.  See
    ``argsort`` for the key dtypes.
    """
    arrays = (keys,) if values is None else (keys, values)
    n = _count(n, *arrays)
    perm = argsort(keys, n)
    sorted_keys = gather(keys, perm, n)
    sorted_values = None if values is None else gather(values, perm, n)
    return sorted_keys, sorted_values

tack.algorithms.sort.gather

gather(src, indices, n=None)

New array out with out[i] = src[indices[i]] for i < n.

indices is an integer array; n defaults to its size. The output has src's dtype. Indices must lie in [0, src.size).

Source code in packages/tack-core/src/tack/algorithms/sort.py

@on_inputs_device
def gather(src, indices, n=None):
    """New array ``out`` with ``out[i] = src[indices[i]]`` for ``i < n``.

    ``indices`` is an integer array; ``n`` defaults to its size.  The
    output has ``src``'s dtype.  Indices must lie in ``[0, src.size)``.
    """
    n = _count(n, indices)
    out = tack.empty((n,), element_type(src))
    if n:
        _gather(src, indices, out, n)
    return out

tack.algorithms.sort.unique

unique(sorted_keys, n=None)

The distinct keys of a sorted array and how often each occurs.

sorted_keys must have equal keys adjacent (sorted, or grouped). Returns (keys, counts): keys is a new array of the same dtype holding each distinct key once, in order of first appearance, and counts the run lengths, an index array (see argsort). Works for any dtype that supports !=.

Source code in packages/tack-core/src/tack/algorithms/sort.py

@on_inputs_device
def unique(sorted_keys, n=None):
    """The distinct keys of a sorted array and how often each occurs.

    ``sorted_keys`` must have equal keys adjacent (sorted, or grouped).
    Returns ``(keys, counts)``: ``keys`` is a new array of the same dtype
    holding each distinct key once, in order of first appearance, and
    ``counts`` the run lengths, an index array (see ``argsort``).  Works for any
    dtype that supports ``!=``.
    """
    n = _count(n, sorted_keys)
    offsets, nruns = _run_offsets(sorted_keys, n)
    keys = tack.empty((nruns,), element_type(sorted_keys))
    counts = tack.empty((nruns,), offsets.dtype)
    if nruns:
        _gather(sorted_keys, offsets, keys, nruns)
        _run_lengths(offsets, counts, nruns)
    return keys, counts

tack.algorithms.sort.reduce_by_key

reduce_by_key(sorted_keys, values, n=None, op='sum')

Reduce the values of each run of equal keys to one value.

sorted_keys must have equal keys adjacent, as for unique; values is any array with at least n elements. op is "sum", "min" or "max". Returns (keys, reduced): the distinct keys and a new array of values' dtype with one entry per key. Each run is reduced serially in element order, so the result is exact for integers and reproducible for floats; integer sums wrap at the values' width.

Source code in packages/tack-core/src/tack/algorithms/sort.py

@on_inputs_device
def reduce_by_key(sorted_keys, values, n=None, op="sum"):
    """Reduce the values of each run of equal keys to one value.

    ``sorted_keys`` must have equal keys adjacent, as for ``unique``;
    ``values`` is any array with at least ``n`` elements.  ``op`` is
    ``"sum"``, ``"min"`` or ``"max"``.  Returns ``(keys, reduced)``: the
    distinct keys and a new array of ``values``' dtype with one entry per
    key.  Each run is reduced serially in element order, so the result is
    exact for integers and reproducible for floats; integer sums wrap at
    the values' width.
    """
    reducer = _REDUCERS.get(op)
    if reducer is None:
        raise ValueError(f"op must be one of {', '.join(_REDUCERS)}; got {op!r}")
    n = _count(n, sorted_keys, values)
    offsets, nruns = _run_offsets(sorted_keys, n)
    keys = tack.empty((nruns,), element_type(sorted_keys))
    reduced = tack.empty((nruns,), element_type(values))
    if nruns:
        _gather(sorted_keys, offsets, keys, nruns)
        reducer(values, offsets, reduced, nruns)
    return keys, reduced

Copy and fill

copy and fill_value are re-exported from tack.algorithms; copy_with_offset is imported from tack.algorithms.copy.

tack.algorithms.copy.copy

copy(src, dst, n)

Copy n elements from src array to dst array with a kernel.

Source code in packages/tack-core/src/tack/algorithms/copy.py

def copy(src, dst, n):
    """Copy n elements from src array to dst array with a kernel."""
    _copy_kernel(src, dst, n)

tack.algorithms.copy.fill_value

fill_value(dst, value, n)

Fill n elements of dst array with a scalar value with a kernel.

Source code in packages/tack-core/src/tack/algorithms/copy.py

def fill_value(dst, value, n):
    """Fill n elements of dst array with a scalar value with a kernel."""
    _fill_kernel(dst, value, n)

tack.algorithms.copy.copy_with_offset

copy_with_offset(src, dst, dst_offset, n)

Copy n elements from src into dst starting at dst_offset with a kernel.

Source code in packages/tack-core/src/tack/algorithms/copy.py

def copy_with_offset(src, dst, dst_offset, n):
    """Copy n elements from src into dst starting at dst_offset with a kernel."""
    _copy_offset_kernel(src, dst, dst_offset, n)

Block reductions

Usable only inside a @tack.kernel, on backends with workgroups (not CPU).

tack.block_sum

block_sum(value)

Sum a value across all threads in the workgroup. Only usable inside a @tack.kernel.

Source code in packages/tack-core/src/tack/__init__.py

def block_sum(value):
    """Sum a value across all threads in the workgroup. Only usable inside a @tack.kernel."""
    raise RuntimeError("block_sum() can only be used inside a @tack.kernel")

tack.block_min

block_min(value)

Min of a value across all threads in the workgroup. Only usable inside a @tack.kernel.

Source code in packages/tack-core/src/tack/__init__.py

def block_min(value):
    """Min of a value across all threads in the workgroup. Only usable inside a @tack.kernel."""
    raise RuntimeError("block_min() can only be used inside a @tack.kernel")

tack.block_max

block_max(value)

Max of a value across all threads in the workgroup. Only usable inside a @tack.kernel.

Source code in packages/tack-core/src/tack/__init__.py

def block_max(value):
    """Max of a value across all threads in the workgroup. Only usable inside a @tack.kernel."""
    raise RuntimeError("block_max() can only be used inside a @tack.kernel")

Communication between blocks

tack.communication moves data between blocks -- units of work with global ids, each on a device, such as the pieces of a partitioned dataset -- and combines values over them. World() is this process; its blocks may be on any of its devices. The Visualization Guide's Communicate Between Blocks builds an algorithm on it step by step, and the Runtime API contract says what it guarantees.

tack.communication.World

World(comm=None, *, gpu_aware=False)

The processes blocks live in, and how messages move between them.

World() is this process alone (rank 0 of 1): its blocks may be on any of its devices, and a message reaches its destination as copies of its arrays on the destination block's device -- the arrays themselves when that device can already address them.

World(comm) spans the ranks of comm, an mpi4py communicator (or anything with its Dup, Get_rank, Get_size, allgather, alltoall, Isend and Irecv). Tack communicates on a duplicate, so its messages never meet the application's; make one World for a run. Messages between ranks are staged through host memory, or with gpu_aware=True passed to MPI as the arrays themselves (through DLPack), for an MPI that reads and writes device memory. Where a backend says the MPI may not take an array's memory as it is (the stream-ordered pools of CUDA and HIP, which Cray MPICH cannot export and MVAPICH-Plus's UCX takes for host memory), the array goes as a device copy in memory the MPI does take, and arrives in such memory. Tack imports no MPI module: the caller makes the communicator.

A GPU-aware MPI uses the device that is current when MPI starts. Make the rank's device first, with tack.device("cuda:N"), and only then import mpi4py.MPI, which starts MPI. A rank that moves to another device afterwards loses Cray MPICH's CUDA IPC between ranks on a node (7 GB/s a rank against 22 on Perlmutter's A100s), or fails across nodes. CUDA IPC also needs every rank to see its node's GPUs (Slurm's --gpu-bind=none, not --gpus-per-task=1), each rank taking its own by local rank.

Source code in packages/tack-core/src/tack/communication/_world.py

def __init__(self, comm=None, *, gpu_aware=False):
    self.gpu_aware = bool(gpu_aware)
    if comm is None:
        if gpu_aware:
            raise ValueError("gpu_aware applies to messages between ranks; "
                             "World() is one process")
        self._comm = None
        self.rank, self.size = 0, 1
    else:
        self._comm = comm.Dup()
        self.rank, self.size = self._comm.Get_rank(), self._comm.Get_size()

tack.communication.Blocks

Blocks(devices, world=None)

This rank's blocks, one per entry of devices (a tack.Device or a spec such as "cuda:1" each), in a World (World(), this process alone, if not given).

Global ids run over every rank's blocks, contiguous by rank: rank 0's first, then rank 1's, as VTK numbers a partitioned dataset's pieces across ranks. local is the range of this rank's.

Source code in packages/tack-core/src/tack/communication/_blocks.py

def __init__(self, devices, world=None):
    self.world = World() if world is None else world
    self._devices = [backend_of(d) for d in devices]
    counts = self.world._allgather(len(self._devices))
    starts = np.concatenate(([0], np.cumsum(counts))).astype(int)
    self._starts = starts
    self.count = int(starts[-1])
    self.local = range(int(starts[self.world.rank]), int(starts[self.world.rank + 1]))

rank_of

rank_of(gid)

The rank holding block gid, any block's: global ids run over rank 0's blocks, then rank 1's, and so on.

Source code in packages/tack-core/src/tack/communication/_blocks.py

def rank_of(self, gid):
    """The rank holding block ``gid``, any block's: global ids run over rank 0's
    blocks, then rank 1's, and so on."""
    self._check_gid(gid, "block")
    return int(np.searchsorted(self._starts, gid, side="right")) - 1

device

device(gid)

The tack.Device of block gid, one of this rank's: where its arrays are, sent and received. ValueError for another rank's block.

Source code in packages/tack-core/src/tack/communication/_blocks.py

def device(self, gid):
    """The `tack.Device` of block ``gid``, one of this rank's: where its arrays
    are, sent and received. ``ValueError`` for another rank's block."""
    return as_device(self._backend(gid))

exchange

exchange(outgoing)

Send messages between blocks and return what each block received.

outgoing is {source: {destination: message}}: for blocks on this rank, a message, a dict of names to Tack arrays on the source block's device, for any block. The result is {destination: {source: message}} for every block on this rank, in order of global id, with an empty dict where nothing came; each message's arrays are on the destination block's device.

A message's arrays belong to the receiver afterwards: where the destination's device addresses the source's memory, they are the sent arrays themselves, so a sender must not write to what it sent.

Source code in packages/tack-core/src/tack/communication/_blocks.py

def exchange(self, outgoing):
    """Send messages between blocks and return what each block received.

    ``outgoing`` is ``{source: {destination: message}}``: for blocks on this
    rank, a message, a dict of names to Tack arrays on the source block's
    device, for any block. The result is ``{destination: {source: message}}``
    for every block on this rank, in order of global id, with an empty dict
    where nothing came; each message's arrays are on the destination block's
    device.

    A message's arrays belong to the receiver afterwards: where the
    destination's device addresses the source's memory, they are the sent
    arrays themselves, so a sender must not write to what it sent.
    """
    self._local_keys(outgoing, "exchange()")
    messages = []
    for source in sorted(outgoing):
        device = self._backend(source)
        for destination in sorted(outgoing[source]):
            self._check_gid(destination, "a message's destination")
            message = outgoing[source][destination]
            _check_message(message, source, destination, device)
            messages.append((source, destination, message))
    incoming = {gid: {} for gid in self.local}
    for source, destination, message in self.world._exchange(messages, self.rank_of, self.device):
        incoming[destination][source] = message
    return {gid: dict(sorted(received.items())) for gid, received in incoming.items()}

all_gather

all_gather(values)

Every block's value, a list indexed by global id, on every rank.

values is {gid: value} for every block on this rank; the values are small Python objects (over MPI they are pickled).

Source code in packages/tack-core/src/tack/communication/_blocks.py

def all_gather(self, values):
    """Every block's value, a list indexed by global id, on every rank.

    ``values`` is ``{gid: value}`` for every block on this rank; the values
    are small Python objects (over MPI they are pickled)."""
    self._every_local(values, "all_gather()")
    gathered = {}
    for part in self.world._allgather({gid: values[gid] for gid in self.local}):
        gathered.update(part)
    return [gathered[gid] for gid in range(self.count)]

all_reduce

all_reduce(values, op='sum')

values ({gid: number or NumPy array}, every block on this rank) combined over every block by op ("sum", "min" or "max").

Values combine in order of global id, the same order whichever rank holds which block, so a floating-point sum has the same bits however the blocks are assigned.

Source code in packages/tack-core/src/tack/communication/_blocks.py

def all_reduce(self, values, op="sum"):
    """``values`` (``{gid: number or NumPy array}``, every block on this rank)
    combined over every block by ``op`` (``"sum"``, ``"min"`` or ``"max"``).

    Values combine in order of global id, the same order whichever rank
    holds which block, so a floating-point sum has the same bits however
    the blocks are assigned."""
    if op not in _OPS:
        raise ValueError(f"all_reduce() takes op 'sum', 'min' or 'max', not {op!r}")
    every = self.all_gather(values)
    if not every:
        raise ValueError("all_reduce() over no blocks has no value")
    scalar, array = _OPS[op]
    result = every[0]
    for value in every[1:]:
        both_numbers = not isinstance(result, np.ndarray) and \
            not isinstance(value, np.ndarray)
        result = scalar(result, value) if both_numbers else array(result, value)
    return result

exclusive_scan

exclusive_scan(counts)

Each block's offset, the sum of the counts of the blocks before it: ({gid: offset}, total) for the blocks on this rank, from counts ({gid: int}, every block on this rank) -- where each block's share of a global numbering starts.

Source code in packages/tack-core/src/tack/communication/_blocks.py

def exclusive_scan(self, counts):
    """Each block's offset, the sum of the counts of the blocks before it:
    ``({gid: offset}, total)`` for the blocks on this rank, from ``counts``
    (``{gid: int}``, every block on this rank) -- where each block's share of
    a global numbering starts."""
    every = self.all_gather(counts)
    for gid, count in enumerate(every):
        if not isinstance(count, (int, np.integer)) or count < 0:
            raise ValueError(f"exclusive_scan() takes counts that are integers of at "
                             f"least 0; block {gid}'s is {count!r}")
    starts = np.concatenate(([0], np.cumsum(every, dtype=np.int64)))
    return {gid: int(starts[gid]) for gid in self.local}, int(starts[-1])

link_by_bounds

link_by_bounds(boxes, margin=0.0)

Each block's neighbours: {gid: [other gids]} for the blocks on this rank, every block whose box overlaps or touches its box grown by margin on every side.

boxes is {gid: (lower, upper)} for every block on this rank, two corners of equal length. Blocks that share a face touch, so a margin of 0 finds the neighbours of a decomposition into abutting pieces; a margin of r finds every block within r.

Source code in packages/tack-core/src/tack/communication/_blocks.py

def link_by_bounds(self, boxes, margin=0.0):
    """Each block's neighbours: ``{gid: [other gids]}`` for the blocks on this
    rank, every block whose box overlaps or touches its box grown by
    ``margin`` on every side.

    ``boxes`` is ``{gid: (lower, upper)}`` for every block on this rank, two
    corners of equal length. Blocks that share a face touch, so a margin of
    0 finds the neighbours of a decomposition into abutting pieces; a margin
    of ``r`` finds every block within ``r``."""
    if margin < 0:
        raise ValueError(f"link_by_bounds() takes a margin of at least 0, not {margin}")
    for gid, box in boxes.items():
        lower, upper = (np.asarray(corner, float) for corner in box)
        if lower.shape != upper.shape or lower.ndim != 1:
            raise ValueError(f"block {gid}'s box needs two corners of equal length, "
                             f"not {box!r}")
        if (lower > upper).any():
            raise ValueError(f"block {gid}'s box has a lower corner above its upper "
                             f"one: {box!r}")
    every = self.all_gather({gid: tuple(np.asarray(c, float) for c in box)
                             for gid, box in boxes.items()})
    dims = {len(lower) for lower, _ in every}
    if len(dims) > 1:
        raise ValueError(f"link_by_bounds() needs boxes of one dimension, not {sorted(dims)}")
    lower = np.array([box[0] for box in every]).reshape(self.count, -1)
    upper = np.array([box[1] for box in every]).reshape(self.count, -1)
    links = {}
    for gid in self.local:
        near = np.all((lower <= upper[gid] + margin) & (lower[gid] <= upper + margin), axis=1)
        near[gid] = False
        links[gid] = [int(g) for g in np.flatnonzero(near)]
    return links

tack.communication.Plan

Plan(blocks, sends, receives)

Which elements each block sends to which, and where the receiver puts them: an exchange kept for reuse, as ghost values are refreshed from their owners.

sends is {gid: {destination: indices}} and receives {gid: {source: indices}}, for blocks on this rank of blocks: the element indices a block reads for each destination, and those it writes from each source, in matching order. Indices are integer Tack arrays on the block's device, or anything numpy.asarray makes integers of. Every send must have a receive of its length on the other block, and the other way round; this is checked across ranks when the plan is made. A block may send to itself, as a piece in a periodic box is its own neighbour: what it sends is gathered before anything is written.

Source code in packages/tack-core/src/tack/communication/_plan.py

def __init__(self, blocks, sends, receives):
    self.blocks = blocks
    blocks._local_keys(sends, "Plan's sends")
    blocks._local_keys(receives, "Plan's receives")
    self._sends = _indices(blocks, sends, "sends", "destination")
    self._receives = _indices(blocks, receives, "receives", "source")
    counts = blocks.all_gather({gid: (
        {other: idx.size for other, idx in self._sends.get(gid, {}).items()},
        {other: idx.size for other, idx in self._receives.get(gid, {}).items()})
        for gid in blocks.local})
    for gid in blocks.local:
        sent, received = counts[gid]
        for other, count in sent.items():
            theirs = counts[other][1].get(gid, 0)
            if theirs != count:
                raise ValueError(f"block {gid} sends {count} elements to block {other}, "
                                 f"which receives {theirs} from it")
        for other, count in received.items():
            theirs = counts[other][0].get(gid, 0)
            if theirs != count:
                raise ValueError(f"block {gid} receives {count} elements from block "
                                 f"{other}, which sends {theirs} to it")

run

run(values)

Copy each sent element of values into its place on the receiving block.

values is {gid: array} -- or {gid: {name: array}} to move several arrays in one exchange -- with an entry for every block on this rank that sends or receives; each array is on its block's device, and an array and its counterparts on other blocks have one element type. Received elements are written in place; the others are left as they were.

Source code in packages/tack-core/src/tack/communication/_plan.py

def run(self, values):
    """Copy each sent element of ``values`` into its place on the receiving block.

    ``values`` is ``{gid: array}`` -- or ``{gid: {name: array}}`` to move
    several arrays in one exchange -- with an entry for every block on this
    rank that sends or receives; each array is on its block's device, and an
    array and its counterparts on other blocks have one element type. Received
    elements are written in place; the others are left as they were."""
    blocks = self.blocks
    if len({isinstance(entry, dict) for entry in values.values()}) > 1:
        raise TypeError("Plan.run() takes an array for every block or a dict of "
                        "arrays for every block, not both")
    named = {gid: entry if isinstance(entry, dict) else {"values": entry}
             for gid, entry in values.items()}
    blocks._local_keys(named, "Plan.run()")
    needed = sorted(set(self._sends) | set(self._receives))
    missing = [gid for gid in needed if gid not in named]
    if missing:
        raise ValueError(f"Plan.run() needs values for blocks {missing}, which send or "
                         f"receive")
    for gid in needed:
        for name, array in named[gid].items():
            if not isinstance(array, Array):
                raise TypeError(f"Plan.run() moves Tack arrays; block {gid}'s {name!r} "
                                f"is a {type(array).__name__}")
            if array._backend.memory_domain != blocks._backend(gid).memory_domain:
                raise ValueError(f"block {gid}'s {name!r} is on {array.device}, not on "
                                 f"the block's device {blocks.device(gid)}")
    names = {gid: sorted(named[gid]) for gid in needed}
    if len({tuple(n) for n in names.values()}) > 1:
        raise ValueError(f"Plan.run() needs the same names on every block, not "
                         f"{sorted({tuple(n) for n in names.values()})}")

    outgoing = {}
    for gid, destinations in self._sends.items():
        outgoing[gid] = {}
        for destination, indices in destinations.items():
            message = {}
            for name, array in named[gid].items():
                message[name] = gather(array, indices)
            outgoing[gid][destination] = message
    incoming = blocks.exchange(outgoing)
    for gid, sources in self._receives.items():
        for source, indices in sources.items():
            message = incoming[gid][source]
            for name, array in named[gid].items():
                received = message[name]
                if received.dtype != array.dtype:
                    raise TypeError(f"block {source} sent {name!r} as {received.dtype}, "
                                    f"and block {gid} holds {array.dtype}")
                scatter(received, indices, array)

tack.communication.Routing

Routing(blocks, positions, boxes, holds)

Where each point goes: the lowest block that holds it.

positions is {gid: points} for every block on this rank, an (n, d) host array of coordinates each (n may be 0); boxes {gid: (lower, upper)} for every block on this rank, None for a block that holds nothing. A point is asked of every block whose box holds it (with a relative slack of 1e-6), and holds(gid, points) -- called by the block asked, with the points it was asked about as an (m, d) float64 host array -- says how well it holds each: 0 not at all, or a rank (True is 1). The point goes to the block of the highest rank, the lowest such block on a tie: a locator that holds a point beats one only within its tolerance. Coordinates travel as their f64 bits in u64 arrays, which every device stores, so a block decides on exactly the points it was given.

destination is {gid: array}, each point's block, or -1 where no block holds it. send moves values with the points to their destinations, and reply returns values from there to where the points came from.

Source code in packages/tack-core/src/tack/communication/_route.py

def __init__(self, blocks, positions, boxes, holds):
    self.blocks = blocks
    blocks._every_local(positions, "Routing's positions")
    blocks._every_local(boxes, "Routing's boxes")
    self._positions = {g: _points(positions[g]) for g in blocks.local}
    every = blocks.all_gather({g: None if boxes[g] is None else
                               tuple(np.asarray(c, np.float64) for c in boxes[g])
                               for g in blocks.local})

    # 1. Each point to every block whose box holds it.
    asked = {}
    for g in blocks.local:
        x = self._positions[g]
        asked[g] = {}
        for d, box in enumerate(every):
            if box is None or not len(x):
                continue
            lower, upper = box
            slack = 1e-6 * max(float(np.max(upper - lower)) if len(upper) else 0.0, 1.0)
            inside = np.all((x >= lower - slack) & (x <= upper + slack), axis=1)
            rows = np.flatnonzero(inside)
            if len(rows):
                message = _bits(x[rows], blocks.device(g))
                message["slot"] = tack.array(rows.astype(np.int64), tack.i64,
                                             device=blocks.device(g))
                asked[g][d] = message
    questions = blocks.exchange(asked)

    # 2. Each block says which it holds.
    answers = {}
    for g in blocks.local:
        answers[g] = {}
        for origin, message in questions[g].items():
            held = np.asarray(holds(g, _unbits(message))).astype(np.int32).reshape(-1)
            answers[g][origin] = {"slot": message["slot"],
                                  "held": tack.array(held, tack.i32,
                                                     device=blocks.device(g))}
    replies = blocks.exchange(answers)

    # 3. The highest rank takes it, the lowest block on a tie (replies come
    # in order of block, so a later one must rank strictly higher).
    self.destination = {}
    for g in blocks.local:
        n = len(self._positions[g])
        destination = np.full(n, -1, np.int64)
        best = np.zeros(n, np.int32)
        for d, reply in replies[g].items():
            slots = reply["slot"].numpy()
            rank = reply["held"].numpy()
            better = rank > best[slots]
            destination[slots[better]] = d
            best[slots[better]] = rank[better]
        self.destination[g] = destination
    self._arrivals = None

send

send(values=None)

Move values ({gid: {name: array}}, an element per point, on the block's device; may be omitted) with the points to their destinations. Returns {gid: (positions, {name: array})} for every block on this rank: what arrived there, the points as an (m, d) host array and the values on the block's device, in order of origin block and then of the points there. reply answers in that order.

Source code in packages/tack-core/src/tack/communication/_route.py

def send(self, values=None):
    """Move ``values`` (``{gid: {name: array}}``, an element per point, on the
    block's device; may be omitted) with the points to their destinations.
    Returns ``{gid: (positions, {name: array})}`` for every block on this
    rank: what arrived there, the points as an ``(m, d)`` host array and the
    values on the block's device, in order of origin block and then of the
    points there. ``reply`` answers in that order."""
    blocks = self.blocks
    values = values or {g: {} for g in blocks.local}
    outgoing, self._order = {}, {}
    for g in blocks.local:
        outgoing[g] = {}
        for d in np.unique(self.destination[g][self.destination[g] >= 0]).tolist():
            rows = np.flatnonzero(self.destination[g] == d)
            device = blocks.device(g)
            message = _bits(self._positions[g][rows], device)
            message["slot"] = tack.array(rows.astype(np.int64), tack.i64, device=device)
            indices = tack.array(rows.astype(np.int64), tack.i64, device=device)
            for name, array in values.get(g, {}).items():
                if not isinstance(array, Array):
                    raise TypeError(f"Routing.send() moves Tack arrays; block {g}'s "
                                    f"{name!r} is a {type(array).__name__}")
                message["value." + name] = gather(array, indices)
            outgoing[g][int(d)] = message
    incoming = blocks.exchange(outgoing)
    arrived = {}
    self._arrivals = {}
    for g in blocks.local:
        parts = list(incoming[g].items())               # in order of origin
        self._arrivals[g] = [(origin, m["slot"], m["slot"].size) for origin, m in parts]
        x = (np.concatenate([_unbits(m) for _, m in parts]) if parts
             else np.zeros((0, self._width()), np.float64))
        names = sorted({k for _, m in parts for k in m if k.startswith("value.")})
        joined = {}
        for name in names:
            pieces = [m[name] for _, m in parts]
            joined[name[len("value."):]] = pieces[0] if len(pieces) == 1 \
                else tack.concat(pieces)
        arrived[g] = (x, joined)
    return arrived

reply

reply(values, like)

Send values -- {gid: {name: array}}, an element for each point that arrived at the block (send's order) -- back to where the points came from. Returns {gid: {name: array}} for the blocks on this rank, an element per point of the block's positions, zero where the point had no destination. like is {name: element type}, for blocks none of whose points were taken.

Source code in packages/tack-core/src/tack/communication/_route.py

def reply(self, values, like):
    """Send ``values`` -- ``{gid: {name: array}}``, an element for each point
    that arrived at the block (``send``'s order) -- back to where the points
    came from. Returns ``{gid: {name: array}}`` for the blocks on this rank,
    an element per point of the block's ``positions``, zero where the point
    had no destination. ``like`` is ``{name: element type}``, for blocks
    none of whose points were taken."""
    if self._arrivals is None:
        raise ValueError("Routing.reply() answers what Routing.send() delivered: "
                         "send first")
    blocks = self.blocks
    outgoing = {}
    for g in blocks.local:
        outgoing[g] = {}
        start = 0
        for origin, slots, count in self._arrivals[g]:
            indices = tack.array(np.arange(start, start + count, dtype=np.int64),
                                 tack.i64, device=blocks.device(g))
            message = {"slot": slots.to_device(blocks.device(g))}
            for name in like:
                message["value." + name] = gather(values[g][name], indices)
            outgoing[g][origin] = message
            start += count
    incoming = blocks.exchange(outgoing)
    out = {}
    for g in blocks.local:
        n = len(self._positions[g])
        out[g] = {}
        for name, dtype in like.items():
            array = tack.zeros((n,), dtype, device=blocks.device(g))
            for _, message in incoming[g].items():
                scatter(message["value." + name], message["slot"], array)
            out[g][name] = array
    return out

tack.communication.sort_by_key

sort_by_key(blocks, keys, values=None, samples=16)

Sort the keys of every block, carrying values along: block 0 ends with the smallest keys, the last block with the largest, each block's ascending.

keys is {gid: array} and values {gid: {name: array}}, for every block on this rank: one-dimensional i32, u32, i64 or u64 keys of one type on every block, and values of the keys' length, each name of one element type everywhere, all on the block's device. Returns (keys, values) in the same form (values None if not given), on the same devices.

Every copy of a key ends on one block, so equal keys can be grouped where they land; heavily repeated keys can leave blocks uneven. The sort is stable: equal keys keep the order of their blocks, then their order within a block. A sample sort, as DIY's: each block sorts its own keys and sends samples of them, weighted by its count, from which every block finds the same count - 1 splitters; each block sends each range to its block, and sorts what arrives.

Source code in packages/tack-core/src/tack/communication/_sort.py

def sort_by_key(blocks, keys, values=None, samples=16):
    """Sort the keys of every block, carrying ``values`` along: block 0 ends with the
    smallest keys, the last block with the largest, each block's ascending.

    ``keys`` is ``{gid: array}`` and ``values`` ``{gid: {name: array}}``, for every
    block on this rank: one-dimensional ``i32``, ``u32``, ``i64`` or ``u64`` keys of
    one type on every block, and values of the keys' length, each name of one
    element type everywhere, all on the block's device. Returns ``(keys,
    values)`` in the same form (``values`` None if not given), on the same
    devices.

    Every copy of a key ends on one block, so equal keys can be grouped where
    they land; heavily repeated keys can leave blocks uneven. The sort is
    stable: equal keys keep the order of their blocks, then their order within
    a block. A sample sort, as DIY's: each block sorts its own keys and sends
    ``samples`` of them, weighted by its count, from which every block finds the
    same ``count - 1`` splitters; each block sends each range to its block, and
    sorts what arrives.
    """
    blocks._every_local(keys, "sort_by_key()")
    given = values is not None
    values = {gid: {} for gid in blocks.local} if values is None else values
    blocks._every_local(values, "sort_by_key()'s values")
    for gid in blocks.local:
        _check(blocks, gid, keys[gid], values[gid])
    kinds = blocks.all_gather({gid: (keys[gid].dtype, {name: v.dtype for name, v in
                                                       values[gid].items()})
                               for gid in blocks.local})
    if len({repr(kind) for kind in kinds}) > 1:
        raise TypeError(f"sort_by_key() needs one key type and the same values, of one "
                        f"type each, on every block, not {sorted({repr(k) for k in kinds})}")

    ordered, carried, picks = {}, {}, {}
    for gid in blocks.local:
        perm = argsort(keys[gid])
        ordered[gid] = gather(keys[gid], perm)
        carried[gid] = {name: gather(v, perm) for name, v in values[gid].items()}
        n = ordered[gid].size
        m = min(samples, n)
        at = ((np.arange(m) + 0.5) * n / max(m, 1)).astype(np.int64)
        chosen = gather(ordered[gid], tack.array(at, tack.i64, device=blocks.device(gid)))
        picks[gid] = (chosen.numpy(), n)
    splitters = _splitters(blocks.all_gather(picks), blocks.count)
    if _in_one_place(blocks):
        return _sorted_in_one_place(blocks, ordered, carried, splitters, kinds, given)

    outgoing = {}
    for gid in blocks.local:
        n = ordered[gid].size
        bounds = [0, *search_right(ordered[gid], splitters).tolist(), n]
        outgoing[gid] = {}
        for destination in range(blocks.count):
            start, count = bounds[destination], bounds[destination + 1] - bounds[destination]
            if count:
                message = {"keys": take(ordered[gid], start, count)}
                for name, v in carried[gid].items():
                    message["values." + name] = take(v, start, count)
                outgoing[gid][destination] = message
    incoming = blocks.exchange(outgoing)

    sorted_keys, sorted_values = {}, {}
    for gid in blocks.local:
        runs = list(incoming[gid].values())             # in order of source block
        device = blocks.device(gid)
        key_type, value_types = kinds[gid]
        joined = _joined([run["keys"] for run in runs], key_type, device)
        perm = argsort(joined)
        sorted_keys[gid] = gather(joined, perm)
        sorted_values[gid] = {
            name: gather(_joined([run["values." + name] for run in runs], dtype, device), perm)
            for name, dtype in value_types.items()}
    return sorted_keys, (sorted_values if given else None)

VTK interop

tack.interop.vtk (in the tack-vis package) provides vtk_to_array, array_to_vtk and init_level_zero, and vtk_to_dataset/dataset_to_vtk. Their contracts are in Runtime API: VTK, and their entries, with MFEM's and Warp's, in the Visualization Guide's Data Model. The tack.dataset algorithms are in Algorithms.

The backend contract

Every backend subclasses this. It exists so callers can ask what a backend supports instead of probing for methods — the distinction that hasattr(backend, 'reduce_array') got wrong, since it answers "is this defined" rather than "is this supported".

tack.runtime.backend.Backend

Base class for Tack compute backends.

Subclasses declare their capabilities as class attributes (or set them in __init__ when they are device-dependent, as Level Zero does for f64) and implement the methods below.

supported_atomic_dtypes property

supported_atomic_dtypes: frozenset[ScalarType]

Global array types supporting add/min/max on this target.

supports_f64 property

supports_f64: bool

Whether kernels can be dispatched with f64 arrays.

Derived, so it cannot disagree with supported_dtypes.

memory_domain property

memory_domain: tuple

What an array's memory is valid in: two backends with one domain can use each other's arrays (two CPU backends with different thread counts, say). CUDA and Level Zero add their context, since a pointer belongs to one.

label property

label: str

The backend's name as it should appear in a message.

device_count classmethod

device_count() -> int

How many devices of this type the machine has.

Source code in packages/tack-core/src/tack/runtime/backend.py

@classmethod
def device_count(cls) -> int:
    """How many devices of this type the machine has."""
    return 1

current_ordinal classmethod

current_ordinal() -> int

The device a bare spec names ("cuda"): the one the calling thread already uses, as another library may have chosen it.

Source code in packages/tack-core/src/tack/runtime/backend.py

@classmethod
def current_ordinal(cls) -> int:
    """The device a bare spec names (``"cuda"``): the one the calling
    thread already uses, as another library may have chosen it."""
    return 0

allocate_array

allocate_array(
    dtype: ScalarType,
    shape: tuple[int, ...],
    exportable: bool = False,
)

Allocate backend storage for an array, returning a DeviceBuffer.

Source code in packages/tack-core/src/tack/runtime/backend.py

def allocate_array(self, dtype: ScalarType, shape: tuple[int, ...],
                   exportable: bool = False):
    """Allocate backend storage for an array, returning a DeviceBuffer."""
    raise NotImplementedError

wrap_ptr

wrap_ptr(ptr, dtype: ScalarType, shape: tuple[int, ...])

Wrap existing memory as a DeviceBuffer without copying.

Source code in packages/tack-core/src/tack/runtime/backend.py

def wrap_ptr(self, ptr, dtype: ScalarType, shape: tuple[int, ...]):
    """Wrap existing memory as a DeviceBuffer without copying."""
    raise NotImplementedError

addresses_dlpack

addresses_dlpack(device_type: int, device_id: int) -> bool

Whether this backend can wrap a DLPack tensor on (device_type, device_id) without copying, so from_dlpack may put it here. By default by the backend's name (lang/dlpack.py's _DEVICE_BACKENDS) and, for a type with several devices, its ordinal. The Warp backend answers for the one device it drives.

Source code in packages/tack-core/src/tack/runtime/backend.py

def addresses_dlpack(self, device_type: int, device_id: int) -> bool:
    """Whether this backend can wrap a DLPack tensor on ``(device_type,
    device_id)`` without copying, so `from_dlpack` may put it here. By
    default by the backend's name (``lang/dlpack.py``'s
    ``_DEVICE_BACKENDS``) and, for a type with several devices, its ordinal.
    The Warp backend answers for the one device it drives."""
    from tack.lang.dlpack import _DEVICE_BACKENDS
    return self.name in _DEVICE_BACKENDS.get(device_type, ()) and (
        not self.multi_device or self.ordinal == device_id)

mpi_direct

mpi_direct(buffer) -> bool

Whether a GPU-aware MPI may read and write buffer's memory as it is (tack.communication.World(comm, gpu_aware=True)). True by default; a backend says no for memory such an MPI does not know -- the stream-ordered pools of CUDA and HIP, which Cray MPICH cannot export and MVAPICH-Plus's UCX takes for host memory -- and the exchange goes through allocate_mpi_array's memory instead.

Source code in packages/tack-core/src/tack/runtime/backend.py

def mpi_direct(self, buffer) -> bool:
    """Whether a GPU-aware MPI may read and write ``buffer``'s memory as it is
    (``tack.communication.World(comm, gpu_aware=True)``). True by default;
    a backend says no for memory such an MPI does not know -- the
    stream-ordered pools of CUDA and HIP, which Cray MPICH cannot export
    and MVAPICH-Plus's UCX takes for host memory -- and the exchange goes
    through `allocate_mpi_array`'s memory instead."""
    return True

allocate_mpi_array

allocate_mpi_array(
    dtype: ScalarType, shape: tuple[int, ...]
)

Storage a GPU-aware MPI reads and writes directly (mpi_direct): by default what allocate_array gives.

Source code in packages/tack-core/src/tack/runtime/backend.py

def allocate_mpi_array(self, dtype: ScalarType, shape: tuple[int, ...]):
    """Storage a GPU-aware MPI reads and writes directly (`mpi_direct`): by
    default what `allocate_array` gives."""
    return self.allocate_array(dtype, shape)

fill_array

fill_array(array, value) -> bool

Fill an array with one scalar the backend's own way; false to let Array.fill fill the buffer in place.

The CPU backend runs a large fill as a kernel through its dispatch, so the threads first-touch the pages. GPU buffers fill on the device through DeviceBuffer.fill and need nothing here.

Source code in packages/tack-core/src/tack/runtime/backend.py

def fill_array(self, array, value) -> bool:
    """Fill an array with one scalar the backend's own way; false to let
    `Array.fill` fill the buffer in place.

    The CPU backend runs a large fill as a kernel through its dispatch,
    so the threads first-touch the pages. GPU buffers fill on the
    device through `DeviceBuffer.fill` and need nothing here.
    """
    return False

execute

execute(kernel, args, kwargs)

Compile (or reuse) and run a kernel with these arguments.

Source code in packages/tack-core/src/tack/runtime/backend.py

def execute(self, kernel, args, kwargs):
    """Compile (or reuse) and run a kernel with these arguments."""
    raise NotImplementedError

release_memory

release_memory()

Free memory kept for reuse (a cache of dropped arrays' buffers), so an allocation that ran out of memory can try again (lang.array._allocate). Nothing by default.

Source code in packages/tack-core/src/tack/runtime/backend.py

def release_memory(self):
    """Free memory kept for reuse (a cache of dropped arrays' buffers), so an
    allocation that ran out of memory can try again (``lang.array._allocate``).
    Nothing by default."""

synchronize

synchronize()

Wait until every kernel launched so far has finished.

Backends whose launches complete before execute returns need not override it. Metal queues launches and waits only when the host touches array memory, so tack.sync() is how a caller timing kernels, or handing memory to a library Tack cannot see, waits.

Source code in packages/tack-core/src/tack/runtime/backend.py

def synchronize(self):
    """Wait until every kernel launched so far has finished.

    Backends whose launches complete before ``execute`` returns need not
    override it. Metal queues launches and waits only when the host
    touches array memory, so ``tack.sync()`` is how a caller timing
    kernels, or handing memory to a library Tack cannot see, waits.
    """

async_marker

async_marker()

A mark after every launch queued so far that another thread can wait on -- an object with a blocking wait(), which raises a launch's error as synchronize would -- or None.

tack.sync_async waits on it in a worker thread so an asyncio loop keeps running. None, the default, means there is no such mark, and the coroutine calls synchronize inline: right for a backend whose launches complete before returning, and still correct, if blocking, for one that queues launches and has not added a mark.

Source code in packages/tack-core/src/tack/runtime/backend.py

def async_marker(self):
    """A mark after every launch queued so far that another thread can
    wait on -- an object with a blocking ``wait()``, which raises a
    launch's error as `synchronize` would -- or None.

    `tack.sync_async` waits on it in a worker thread so an ``asyncio``
    loop keeps running. None, the default, means there is no such mark,
    and the coroutine calls `synchronize` inline: right for a backend
    whose launches complete before returning, and still correct, if
    blocking, for one that queues launches and has not added a mark.
    """
    return

memory_space

memory_space(ptr) -> str

Classify where an integer pointer lives.

The default says host memory, which is correct for backends whose allocations are CPU-addressable. Backends with distinct device memory override this and list the results in device_memory_spaces.

Source code in packages/tack-core/src/tack/runtime/backend.py

def memory_space(self, ptr) -> str:
    """Classify where an integer pointer lives.

    The default says host memory, which is correct for backends whose
    allocations are CPU-addressable. Backends with distinct device
    memory override this and list the results in
    `device_memory_spaces`.
    """
    return "cpu"

cache_target

cache_target()

What compiled code for this backend depends on besides the kernel, as a string (the device, the compiler's version), or None if this backend does not store compiled variants on disk. See disk_cache.

Source code in packages/tack-core/src/tack/runtime/backend.py

def cache_target(self):
    """What compiled code for this backend depends on besides the kernel, as a
    string (the device, the compiler's version), or None if this backend
    does not store compiled variants on disk. See ``disk_cache``."""

variant_artifact

variant_artifact(payload)

What of a compiled variant's payload to store on disk: picklable data from which load_variant rebuilds it in another process.

Source code in packages/tack-core/src/tack/runtime/backend.py

def variant_artifact(self, payload):
    """What of a compiled variant's payload to store on disk: picklable data
    from which ``load_variant`` rebuilds it in another process."""
    raise NotImplementedError

load_variant

load_variant(ir_func, artifact, effective_args)

The payload of a variant, from what variant_artifact stored.

effective_args are the call's arguments, for what the payload holds per process rather than on disk: the GPU backends' scalar pack buffers. CPU and Metal ignore them.

Source code in packages/tack-core/src/tack/runtime/backend.py

def load_variant(self, ir_func, artifact, effective_args):
    """The payload of a variant, from what ``variant_artifact`` stored.

    ``effective_args`` are the call's arguments, for what the payload
    holds per process rather than on disk: the GPU backends' scalar
    pack buffers. CPU and Metal ignore them."""
    raise NotImplementedError

reduce_array

reduce_array(array, op: str) -> float

Reduce an array on the device. Only called when supports_device_reductions is True.

Source code in packages/tack-core/src/tack/runtime/backend.py

def reduce_array(self, array, op: str) -> float:
    """Reduce an array on the device. Only called when
    `supports_device_reductions` is True."""
    raise NotImplementedError(
        f"{self.label} backend does not implement device reductions")

texture_in_hardware

texture_in_hardware(shape_3d) -> bool

Whether a Texture3D of this extent is sampled by texture units.

False means the texture keeps a private array and the generated code interpolates it in software, which is what CPU does. Texture3D asks when choosing its storage, and backends whose answer varies by device or extent ask again when preparing a variant, so a texture always holds the storage its compiled kernel binds.

Source code in packages/tack-core/src/tack/runtime/backend.py

def texture_in_hardware(self, shape_3d) -> bool:
    """Whether a `Texture3D` of this extent is sampled by texture units.

    False means the texture keeps a private array and the generated
    code interpolates it in software, which is what CPU does. `Texture3D`
    asks when choosing its storage, and backends whose answer varies by
    device or extent ask again when preparing a variant, so a texture
    always holds the storage its compiled kernel binds.
    """
    return False

create_texture_image

create_texture_image(shape_3d)

Allocate a single-channel f32 3D texture image.

Only called when texture_in_hardware() is True. The result is owned by one Texture3D: it exposes upload(array), which copies the array into it, and frees the device objects when collected.

Source code in packages/tack-core/src/tack/runtime/backend.py

def create_texture_image(self, shape_3d):
    """Allocate a single-channel f32 3D texture image.

    Only called when `texture_in_hardware()` is True. The result is
    owned by one `Texture3D`: it exposes ``upload(array)``, which copies
    the array into it, and frees the device objects when collected.
    """
    raise NotImplementedError(
        f"{self.label} backend has no hardware textures")