i

Python front-end for πš’.

This package exposes πš’ components as Python objects. Components compile lazily, execute on CPU or CUDA according to their inputs, and interoperate with i.Tensor, Python scalar/list literals, NumPy arrays, and Torch tensors.

Package-style API

import ilang

Exports:

Preferred β€œDSL-style” API

The package-exported object i acts as a callable β€œnamespace” that enables a more compact style of πš’ code. When called, it constructs a Component, but it also re-exposes much of the same package-level API as attributes.

from ilang import i

i("+i~.")   # <ilang.component.Component object at ...>
i.Tensor    # <class 'ilang.tensor.Tensor'>
i.Component # <class 'ilang.component.Component'>
i.Device    # <enum 'Device'>
i.I         # mirrors `ilang.Component.I`

Devices

A Device dictates where data lives and where compuation will run.

i.Device.CPU  # "cpu"
i.Device.CUDA # "cuda"

Tensors

Tensors are immutable multidimensional data arrays.

x = i.Tensor([[1, 2], [3, 4]]) # standard construction does shape-inference on nested list
x = i.Tensor([1, 2, 3, 4], shape=(2, 2)) # flat data with shape also works

Values are stored as float32.

Attributes

x.shape  # tuple[int, ...]
x.device # i.Device.CPU or i.Device.CUDA
x.data   # list[float] (only available on CPU tensors)

Methods

Tensor.to(device i.Device) -> Tensor # gives a new tensor on specified device

# examples:
x.to(i.Device.CUDA)
x.to(i.Device.CPU)
x.to("cuda")
x.to("cpu")

Components

i(expr: str) -> Component parses one πš’ expression.

f = i("+ij~i") # row-sum

i.I is the identity component.

Component combinators

Components are combined using combinators. Components are immutable, so combinators each return a new component.

Method forms:

f.compose(g) # wires outputs of the right component into inputs of the left component
f.chain(g)   # wires outputs of the left component into inputs of the right component
f.fanout(g)  # shares inputs pairwise between two components
f.pair(g)    # concatenates the inputs and outputs of two components
f.swap()     # swaps the first two outputs of one component

Operator forms:

f << g # f.compose(g)
f >> g # f.chain(g)
f & g  # f.fanout(g)
f | g  # f.pair(g)
~f     # f.swap()

Example:

matmul = i("ik*kj~ijk") >> i("+ijk~ij")

Execution

out = matmul.exec(x, y)

Component.exec(*inputs: TensorLike, into=None) executes one component.

where TensorLike = Tensor | torch.Tensor | numpy.ndarray | nested Python sequence

Execution device is determined by the input devices. All inputs must use the same device. NumPy/Torch inputs must be contiguous and of dtype float32. Python scalars/lists get promoted to CPU ilang.Tensor.

Output type inference

The output container type and device are inferred from the input. For example, f.exec(torch.tensor([...], device="cuda")) returns an torch.Tensor with device="cuda". This can be overridden with into=: f.exec(nparray, into=i.Tensor).

In general, component execution returns Tuple[Tensor] but is automatically unpacked for single output results.

Shape inference

Component.output_shapes(*inputs) returns one shape tuple per output. Shapes are computed from input shapes without executing any kernels.

Benchmarking

from ilang.testing import bench

result = bench(lambda: f.exec(x), n_warmups=10, n_runs=100)
# or benchmark a fully-bound component:
result = bench(f(x), n_warmups=10, n_runs=100)

bench executes warmup runs, records timed runs, and returns Bench.

Bench fields:

bench.mean       # datetime.timedelta
bench.std        # datetime.timedelta
bench.n_warmups  # int
bench.n_runs     # int
bench.runs       # list[datetime.timedelta]

repr(bench) prints a compact human-readable timing summary.

Errors

Invalid programs, invalid input types, invalid dtypes, non-contiguous arrays, device mismatches, and backend failures raise Python exceptions.

Execution requires all inputs to reside on one device. No implicit CPU/CUDA input synchronization is performed by exec.

CUDA tensor data is not read by repr and is not exposed by .data. Copy to CPU explicitly.

Examples

Native tensor execution:

from ilang import i

f = i.i("+ij~ij")
x = i.Tensor([[1, 2], [3, 4]])
y = f.exec(x)

CUDA tensor execution:

x = i.Tensor([[1, 2], [3, 4]]).to("cuda")
y = f.exec(x)
z = y.to("cpu")

NumPy execution:

import numpy as np

x = np.ones((2, 2), dtype=np.float32)
y = f.exec(x)

Torch CUDA execution:

import torch

x = torch.ones((2, 2), dtype=torch.float32, device="cuda")
y = f.exec(x)