API Reference

Reference documentation is generated from Python docstrings in the grumpy package. For narrative tutorials, start at Home and follow the section links at the bottom of each page.

Top-level API

Grumpy: high-performance numerical computing on ragged and nested data.

Grumpy provides Awkward-like layout semantics with strong typing, explicit nullability, mutable arrays, Zarr-backed lazy I/O via :func:open, and optional compilation of transforms.

Layouts

Arrays use either list-chains (ListOffset -> … -> Leaf) or **UnionScalarList``** (mixed scalar/list rows at one axis). Both are constructed with :func:array`, persisted to Zarr, streamed, and used in dataframes.

Notes

• :func:open exposes lazy on-disk dataframes; row/schema indexing materializes subsets; column access returns lazy column proxies until indexed.
• gr.compile accepts a restricted subset of Python (see :func:compile); scalar elementwise opcodes fuse on union batches as well as list-chains.

bool_ module-attribute

bool_ = DType.bool_()

char module-attribute

char = DType.char()

float16 module-attribute

float16 = DType.float16()

float32 module-attribute

float32 = DType.float32()

float64 module-attribute

float64 = DType.float64()

int16 module-attribute

int16 = DType.int16()

int32 module-attribute

int32 = DType.int32()

int64 module-attribute

int64 = DType.int64()

int8 module-attribute

int8 = DType.int8()

string module-attribute

string = DType.string()

uint16 module-attribute

uint16 = DType.uint16()

uint32 module-attribute

uint32 = DType.uint32()

uint64 module-attribute

uint64 = DType.uint64()

uint8 module-attribute

uint8 = DType.uint8()

abs

abs(x: GrumpyArray) -> GrumpyArray

add

add(a: GrumpyArray, b: GrumpyArray, out: GrumpyArray | None = None) -> GrumpyArray

Elementwise add with optional pre-allocated out.

angle

angle(x: GrumpyArray) -> GrumpyArray

argwhere

argwhere(cond: GrumpyArray) -> GrumpyArray

array

array(obj, dtype: DType | None = None) -> GrumpyArray

Construct a GrumpyArray from Python scalars / nested lists or tuples.

Parameters:

Name	Type	Description	Default
obj		Python scalar or nested Python sequences (lists/tuples) of arbitrary depth.	required
dtype	DType | None	Optional explicit dtype. If omitted, dtype is inferred from non-null leaves.	None

bincount

bincount(x: GrumpyArray, weights: GrumpyArray | None = None, minlength: int = 0) -> GrumpyArray

can_cast

can_cast(from_dtype: DType, to_dtype: DType, casting: str = 'safe') -> bool

Return whether from_dtype can be cast to to_dtype under casting.

cat

cat(arrays: list[GrumpyArray], dim: int = 0) -> GrumpyArray

Concatenate arrays along a ragged dimension.

ceil

ceil(x: GrumpyArray) -> GrumpyArray

cos

cos(x: GrumpyArray) -> GrumpyArray

cross

cross(a: GrumpyArray, b: GrumpyArray) -> GrumpyArray

dataframe

dataframe(mapping: dict, schema=None)

det

det(a: GrumpyArray)

digitize

digitize(x: GrumpyArray, bins: GrumpyArray, right: bool = False) -> GrumpyArray

dot

dot(a: GrumpyArray, b: GrumpyArray)

einsum

einsum(subscripts: str, *operands)

equal

equal(a: GrumpyArray, b: GrumpyArray) -> GrumpyArray

exp

exp(x: GrumpyArray) -> GrumpyArray

floor

floor(x: GrumpyArray) -> GrumpyArray

from_numpy

from_numpy(obj, dtype: DType | None = None) -> GrumpyArray

Build a rectangular :class:GrumpyArray from a C-contiguous NumPy array.

from_tensorflow

from_tensorflow(obj, dtype: DType | None = None) -> GrumpyArray

Build a rectangular :class:GrumpyArray from a TensorFlow tensor (via .numpy()).

from_torch

from_torch(obj, dtype: DType | None = None) -> GrumpyArray

Build a rectangular :class:GrumpyArray from a PyTorch tensor (via .numpy()).

full_like

full_like(x: GrumpyArray, fill_value, dtype: DType | None = None) -> GrumpyArray

Create an array with the same ragged structure as x, filled with fill_value.

gpu_available

gpu_available() -> bool

Return True when a GPU backend (Metal or CUDA) is available.

gpu_backend

gpu_backend() -> str | None

Return 'metal', 'cuda', or None if no GPU backend is active.

greater

greater(a: GrumpyArray, b: GrumpyArray) -> GrumpyArray

greater_equal

greater_equal(a: GrumpyArray, b: GrumpyArray) -> GrumpyArray

grid_pool

grid_pool(x: GrumpyArray, grid_size: tuple[int, int, int], *, origin: tuple[float, float, float] | None = None, voxel_size: tuple[float, float, float] | None = None, dim: int = 1) -> GrumpyArray

Voxelize point clouds by counting points per grid cell (occupancy pooling).

Returns (n_groups, nx*ny*nz) occupancy grids per group.

histogram

histogram(x: GrumpyArray, bins: int = 10, range: tuple[float, float] | None = None, density: bool = False, weights: GrumpyArray | None = None) -> tuple[GrumpyArray, GrumpyArray]

inner

inner(a: GrumpyArray, b: GrumpyArray)

inv

inv(a: GrumpyArray) -> GrumpyArray

is_rectangular

is_rectangular(arr: GrumpyArray) -> bool

Return whether arr can be exported as a dense rectangular tensor.

isfinite

isfinite(x: GrumpyArray) -> GrumpyArray

isin

isin(x: GrumpyArray, test_elements: GrumpyArray) -> GrumpyArray

isinf

isinf(x: GrumpyArray) -> GrumpyArray

isnan

isnan(x: GrumpyArray) -> GrumpyArray

less

less(a: GrumpyArray, b: GrumpyArray) -> GrumpyArray

less_equal

less_equal(a: GrumpyArray, b: GrumpyArray) -> GrumpyArray

load

load(path: str)

Load a saved dataset fully into memory.

log

log(x: GrumpyArray) -> GrumpyArray

log10

log10(x: GrumpyArray) -> GrumpyArray

log2

log2(x: GrumpyArray) -> GrumpyArray

logical_and

logical_and(a: GrumpyArray, b: GrumpyArray) -> GrumpyArray

logical_not

logical_not(a: GrumpyArray) -> GrumpyArray

logical_or

logical_or(a: GrumpyArray, b: GrumpyArray) -> GrumpyArray

logical_xor

logical_xor(a: GrumpyArray, b: GrumpyArray) -> GrumpyArray

median

median(x: GrumpyArray, dim: int = 0) -> GrumpyArray

multiply

multiply(a: GrumpyArray, b: GrumpyArray, out: GrumpyArray | None = None) -> GrumpyArray

Elementwise multiply with optional pre-allocated out (NumPy out= style).

nanmedian

nanmedian(x: GrumpyArray, dim: int = 0) -> GrumpyArray

nanpercentile

nanpercentile(x: GrumpyArray, q: float, dim: int = 0) -> GrumpyArray

nanquantile

nanquantile(x: GrumpyArray, q: float, dim: int = 0) -> GrumpyArray

nanstd

nanstd(x: GrumpyArray, dim: int = 0, ddof: int = 0) -> GrumpyArray

nanvar

nanvar(x: GrumpyArray, dim: int = 0, ddof: int = 0) -> GrumpyArray

neighbors

neighbors(query: GrumpyArray, data: GrumpyArray, k: int | None = None, radius: float | None = None, dim: int = 0, loop: bool = True, return_distances: bool = False, gpu: bool | str = False)

Compute neighbors and return an edge index (and optionally distances).

Parameters:

Name	Type	Description	Default
query	GrumpyArray	Coordinate arrays. For dim=1, shape is (n_groups, n_points, d).	required
data	GrumpyArray	Coordinate arrays. For dim=1, shape is (n_groups, n_points, d).	required
k	int | None	Number of nearest neighbors (mutually exclusive with radius).	None
radius	float | None	Include all neighbors within this distance (mutually exclusive with k).	None
dim	int	Axis along which groups of points live (0 or 1 for grouped clouds).	0
loop	bool	If True, include self-matches where query and data share the same point index.	True
return_distances	bool	If True, also return distances aligned with the neighbor axis.	False
gpu	bool | str	False (default) uses CPU. True always uses GPU when available (Metal on macOS, CUDA on Linux with --features cuda). 'auto' selects GPU only for large enough batches. Use :func:gpu_available to check runtime support.	False

Returns:

Name	Type	Description
edge_index		Ragged edge index with last axis length 2: [src, dst].
	distances (optional):	Returned when return_distances=True.

nonzero

nonzero(x: GrumpyArray) -> GrumpyArray

norm

norm(a: GrumpyArray)

not_equal

not_equal(a: GrumpyArray, b: GrumpyArray) -> GrumpyArray

ones_like

ones_like(x: GrumpyArray, dtype: DType | None = None) -> GrumpyArray

Create an array with the same ragged structure as x, filled with ones.

open

open(path: str, *, cache: str = 'chunks', chunk_budget_mb: int = 256) -> OpenDataFrame

Open a saved dataset as a lazy on-disk handle.

Row and schema indexing materialize subset dataframes. Column selection (df['col'] or df.col) returns an :class:OpenColumn proxy until the column is indexed.

cache controls session I/O caching: "chunks" (default) keeps a byte-bounded LRU of decoded Zarr chunks; "metadata" pins only offset buffers; "none" disables caching.

Call :meth:~grumpy.OpenDataFrame.close when finished, or use a context manager::

with gr.open("data.gr") as h:
    batch = h[[0, 1, 2]]

outer

outer(a: GrumpyArray, b: GrumpyArray) -> GrumpyArray

pairwise_distances

pairwise_distances(x: GrumpyArray, *, dim: int = 1) -> GrumpyArray

All-pairs Euclidean distances within each point cloud (group).

For dim=1, input shape is (n_groups, n_points, d); output is (n_groups, n_points, n_points) distance matrices.

percentile

percentile(x: GrumpyArray, q: float, dim: int = 0) -> GrumpyArray

promote_types

promote_types(a: DType, b: DType) -> DType

NumPy-style binary result dtype for two dtypes.

quantile

quantile(x: GrumpyArray, q: float, dim: int = 0) -> GrumpyArray

reciprocal

reciprocal(x: GrumpyArray) -> GrumpyArray

rng

rng(seed: int = 0) -> Generator

Create a reproducible random :class:~grumpy.Generator.

round

round(x: GrumpyArray) -> GrumpyArray

save

save(obj, path: str, chunk_size: int = 1024, chunk_dim=None)

Save a GrumpyArray/DataFrame, or incrementally write batches from a generator.

search_sorted

search_sorted(x: GrumpyArray, v: GrumpyArray, right: bool = False) -> GrumpyArray

setdiff

setdiff(a: GrumpyArray, b: GrumpyArray) -> GrumpyArray

setunion

setunion(a: GrumpyArray, b: GrumpyArray) -> GrumpyArray

setxor

setxor(a: GrumpyArray, b: GrumpyArray) -> GrumpyArray

sign

sign(x: GrumpyArray) -> GrumpyArray

sin

sin(x: GrumpyArray) -> GrumpyArray

sqrt

sqrt(x: GrumpyArray) -> GrumpyArray

std

std(x: GrumpyArray, dim: int = 0, ddof: int = 0) -> GrumpyArray

subtract

subtract(a: GrumpyArray, b: GrumpyArray, out: GrumpyArray | None = None) -> GrumpyArray

Elementwise subtract with optional pre-allocated out.

tan

tan(x: GrumpyArray) -> GrumpyArray

tensordot

tensordot(a: GrumpyArray, b: GrumpyArray, axes: int = 2)

trace

trace(a: GrumpyArray)

unique

unique(x: GrumpyArray) -> GrumpyArray

var

var(x: GrumpyArray, dim: int = 0, ddof: int = 0) -> GrumpyArray

where

where(cond: GrumpyArray, x: GrumpyArray | None = None, y: GrumpyArray | None = None)

zeros_like

zeros_like(x: GrumpyArray, dtype: DType | None = None) -> GrumpyArray

Create an array with the same ragged structure as x, filled with zeros.

Core types

Compilation

Compile restricted batch transforms into fused Rust execution plans.

The compiler analyzes straight-line Python functions (typically def f(batch): ...) and builds :class:~grumpy._core.CompiledPlan opcode lists for the :func:compile decorator and direct invocation on batches.

Supported inputs

• List-chain and UnionScalarList batches for scalar elementwise opcodes (batch * 2, batch + 1, …).
• Reduction and neighbor opcodes when the underlying Rust kernel supports the layout.

Known limitations

• No control flow (if/for/try), no imports, single batch parameter.
• Rust scheduling of multi-batch pipelines is handled outside Grumpy (e.g. Fabric).

CompiledTransform

CompiledTransform(fn: Callable[[Any], Any], result: _CompileResult)

Callable wrapper that runs a compiled Rust plan when possible.

Instances are returned by :func:compile.

Attributes:

Name	Type	Description
is_compiled	bool	Whether a Rust :class:~grumpy._core.CompiledPlan was built.
compile_error	str or None	Compilation failure message when is_compiled is False.

Examples:

>>> import grumpy as gr
>>> @gr.compile
... def scale(batch):
...     return batch * 2
...
>>> scale.is_compiled
True
>>> scale(gr.array([1, 2])).to_list()
[2, 4]

__doc__ instance-attribute

__doc__ = getattr(fn, '__doc__', None)

__name__ instance-attribute

__name__ = getattr(fn, '__name__', 'compiled_transform')

__qualname__ instance-attribute

__qualname__ = getattr(fn, '__qualname__', self.__name__)

compile_error property

compile_error: Optional[str]

Return the compilation error message, or None on success.

Returns:

Type	Description
str or None	Error text when compilation failed.

Examples:

>>> import grumpy as gr
>>> @gr.compile
... def ok(b): return b * 2
...
>>> ok.compile_error is None
True

is_compiled property

is_compiled: bool

Return True when a Rust :class:~grumpy._core.CompiledPlan was built.

Returns:

Type	Description
bool	Compilation success flag.

Examples:

>>> import grumpy as gr
>>> @gr.compile
... def f(b): return b
...
>>> f.is_compiled
True

__call__

__call__(batch)

Run the compiled plan or fall back to the original Python function.

Parameters:

Name	Type	Description	Default
batch	GrumpyArray or GrumpyDataFrame	Input batch.	required

Returns:

Type	Description
GrumpyArray or GrumpyDataFrame	Transformed batch.

Examples:

>>> import grumpy as gr
>>> @gr.compile
... def double(batch):
...     return batch * 2
...
>>> double(gr.array([1, 2])).to_list()
[2, 4]

compile

compile(fn: Callable[[Any], Any]) -> CompiledTransform

Compile a restricted batch transform into a Rust execution plan.

Parameters:

Name	Type	Description	Default
fn	callable	Function fn(batch) -> batch with straight-line Python only.	required

Returns:

Type	Description
CompiledTransform	Callable wrapper that executes the plan when compilation succeeds.

Examples:

>>> import grumpy as gr
>>> @gr.compile
... def scale(batch):
...     batch = batch * 2
...     return batch
...
>>> scale(gr.array([1, 2])).to_list()
[2, 4]

---

Next: Developer — repository layout, implementation notes, and error handling.