Operations Reference#

This page documents all operations supported by NKIPy, organized by category.

Overview#

NKIPy operations are traced and lowered to HLO for compilation by the Neuron Compiler.

Convolution Operations#

Operation	Backend(s)	NumPy Equivalent
`conv2d`	cpu, hlo	—
`conv3d`	cpu, hlo	—

Binary Operations#

Operation	Backend(s)	NumPy Equivalent
`add`	hlo	`np.add`
`bitwise_and`	hlo	`np.bitwise_and`
`bitwise_or`	hlo	`np.bitwise_or`
`bitwise_xor`	hlo	`np.bitwise_xor`
`divide`	hlo	`np.divide`
`equal`	hlo	`np.equal`
`greater`	hlo	`np.greater`
`greater_equal`	hlo	`np.greater_equal`
`less`	hlo	`np.less`
`less_equal`	hlo	`np.less_equal`
`logical_and`	hlo	`np.logical_and`
`logical_or`	hlo	`np.logical_or`
`logical_xor`	hlo	`np.logical_xor`
`maximum`	hlo	`np.maximum`
`minimum`	hlo	`np.minimum`
`multiply`	hlo	`np.multiply`
`not_equal`	hlo	`np.not_equal`
`power`	hlo	`np.power`
`subtract`	hlo	`np.subtract`

Unary Operations#

Operation	Backend(s)	NumPy Equivalent
`abs`	hlo	`np.abs`
`arctan`	hlo	`np.arctan`
`bitwise_not`	hlo	`np.bitwise_not`
`ceil`	hlo	`np.ceil`
`clip`	hlo	`np.clip`
`cos`	hlo	`np.cos`
`exp`	hlo	`np.exp`
`floor`	hlo	`np.floor`
`invert`	hlo	`np.invert`
`log`	hlo	`np.log`
`log1p`	hlo	`np.log1p`
`logical_not`	hlo	`np.logical_not`
`negative`	hlo	`np.negative`
`reciprocal`	hlo	`np.reciprocal`
`rint`	hlo	`np.rint`
`sign`	hlo	`np.sign`
`sin`	hlo	`np.sin`
`sqrt`	hlo	`np.sqrt`
`square`	hlo	`np.square`
`tan`	hlo	`np.tan`
`tanh`	hlo	`np.tanh`
`trunc`	hlo	`np.trunc`

Reduction Operations#

Operation	Backend(s)	NumPy Equivalent
`any`	hlo	`np.any`
`argmax`	hlo	`np.argmax`
`cumsum`	hlo	`np.cumsum`
`max`	hlo	`np.max`
`mean`	hlo	`np.mean`
`min`	hlo	`np.min`
`std`	hlo	`np.std`
`sum`	hlo	`np.sum`
`var`	hlo	`np.var`

Indexing Operations#

Operation	Backend(s)	NumPy Equivalent
`dynamic_update_slice`	hlo	—
`put_along_axis`	hlo	—
`scatter_along_axis`	hlo	—
`scatter_strided`	hlo	—
`static_slice`	hlo	—
`take`	hlo	`np.take`
`take_along_axis`	hlo	`np.take_along_axis`
`where`	hlo	`np.where`

Neural Network Operations#

Operation	Backend(s)	NumPy Equivalent
`rms_norm`	—	—
`softmax`	—	—
`topk`	cpu, hlo	—

Collective Operations#

Operation	Backend(s)	NumPy Equivalent
`all_gather`	cpu, hlo	—
`all_reduce`	cpu, hlo	—
`all_to_all`	cpu, hlo	—
`reduce_scatter`	cpu, hlo	—

Linear Algebra Operations#

Operation	Backend(s)	NumPy Equivalent
`dot`	hlo	`np.dot`
`matmul`	hlo	`np.matmul`
`norm`	hlo	—

Transform Operations#

Operation	Backend(s)	NumPy Equivalent
`astype`	hlo	—
`broadcast_to`	hlo	`np.broadcast_to`
`concatenate`	hlo	`np.concatenate`
`copy`	hlo	`np.copy`
`copyto`	—	`np.copyto`
`expand_dims`	hlo	`np.expand_dims`
`pad`	hlo	`np.pad`
`repeat`	hlo	`np.repeat`
`reshape`	hlo	`np.reshape`
`split`	hlo	`np.split`
`squeeze`	hlo	`np.squeeze`
`transpose`	hlo	`np.transpose`

Creation Operations#

Operation	Backend(s)	NumPy Equivalent
`constant`	cpu, hlo	—
`empty_like`	cpu, hlo	`np.empty_like`
`full`	cpu, hlo	—
`full_like`	cpu, hlo	`np.full_like`
`ones_like`	cpu, hlo	—
`zeros`	cpu, hlo	—
`zeros_like`	cpu, hlo	`np.zeros_like`

API Reference#

NKIPy Operations Module

This module provides a unified interface for tensor operations that dispatch to the appropriate backend (IR or HLO) based on the current tracing context.

class nkipy.core.ops.Op(name: str)[source]#

Bases: object

Simple operation dispatcher.

Usage:

zeros = Op('zeros')

@zeros.impl('hlo')
def _zeros_hlo(shape, dtype):
    # HLO implementation
    ...

@zeros.impl('cpu')
def _zeros_cpu(shape, dtype):
    # CPU implementation
    ...

# Later, during tracing:
result = zeros(shape, dtype)  # Dispatches based on current backend

__init__(name: str)[source]#

impl(backend: str) → Callable[source]#

Decorator to register a backend implementation.

Parameters:: backend – Backend name (‘hlo’, ‘cpu’, or other supported).
Returns:: Decorator function.

nkipy.core.ops.get_backend() → str[source]#

Return the name of the active backend.

Returns ctx.backend_name when a trace context is active, "cpu" otherwise.

nkipy.core.ops.set_source_location(location: Tuple[str, int] | None) → None[source]#

Set the current source location in the trace context.

Parameters:: location – Tuple of (filename, line_number) or None to clear.

Operations Reference

Contents