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.

Categories

• Convolution Operations

• Neural Network Operations

• Linear Algebra Operations

• Creation Operations

• Unary Operations

• Indexing Operations

• Reduction Operations

• Binary Operations

• Transform Operations

• Collective Operations

Convolution Operations

Operation	Backend(s)	NumPy Equivalent
conv2d	cpu	—
conv3d	cpu	—

Neural Network Operations

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

Linear Algebra Operations

Operation	Backend(s)	NumPy Equivalent
dot	—	np.dot
matmul	—	np.matmul
norm	—	—
outer	—	np.outer
trace	—	np.trace

Creation Operations

Operation	Backend(s)	NumPy Equivalent
constant	cpu	—
diag	—	np.diag
empty_like	cpu	np.empty_like
full	cpu	—
full_like	cpu	np.full_like
ones_like	cpu	np.ones_like
tril	—	np.tril
triu	—	np.triu
zeros	cpu	—
zeros_like	cpu	np.zeros_like

Unary Operations

Operation	Backend(s)	NumPy Equivalent
abs	—	np.abs
arctan	—	np.arctan
bitwise_not	—	np.bitwise_not
ceil	—	np.ceil
clip	—	np.clip
cos	—	np.cos
exp	—	np.exp
expm1	—	np.expm1
floor	—	np.floor
invert	—	np.invert
isfinite	—	np.isfinite
isnan	—	np.isnan
log	—	np.log
log1p	—	np.log1p
log2	—	np.log2
logical_not	—	np.logical_not
negative	—	np.negative
reciprocal	—	np.reciprocal
rint	—	np.rint
round_	—	np.around
sign	—	np.sign
sin	—	np.sin
sqrt	—	np.sqrt
square	—	np.square
tan	—	np.tan
tanh	—	np.tanh
trunc	—	np.trunc

Indexing Operations

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

Reduction Operations

Operation	Backend(s)	NumPy Equivalent
any	—	np.any
argmax	—	np.argmax
argmin	—	np.argmin
count_nonzero	—	np.count_nonzero
cumsum	—	np.cumsum
max	—	np.max
mean	—	np.mean
min	—	np.min
prod	—	np.prod
std	—	np.std
sum	—	np.sum
var	—	np.var

Binary Operations

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

Transform Operations

Operation	Backend(s)	NumPy Equivalent
astype	—	—
broadcast_to	—	np.broadcast_to
concatenate	—	np.concatenate
copy	—	np.copy
copyto	—	np.copyto
diff	—	np.diff
expand_dims	—	np.expand_dims
flip	—	np.flip
pad	—	np.pad
repeat	—	np.repeat
reshape	—	np.reshape
roll	—	np.roll
split	—	np.split
squeeze	—	np.squeeze
stack	—	np.stack
swapaxes	—	np.swapaxes
tile	—	np.tile
transpose	—	np.transpose

Collective Operations

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

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)

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

Composed ops (built from other dispatched ops) use composed_impl:

floor_divide = Op('floor_divide')

@floor_divide.composed_impl
def _floor_divide(x, y):
    return floor(divide(x, y))

The composed fallback is used when no backend-specific implementation is registered. Since it calls other ops via dispatch, it works on any backend that has the underlying primitives registered.

__init__(name: str)

impl(backend: str) → Callable

Decorator to register a backend implementation.

Parameters:

backend – Backend name (‘hlo’, ‘cpu’, or other supported).

Returns:

Decorator function.

property composed_impl

Decorator to register a composed (backend-agnostic) fallback.

A composed implementation is built entirely from calls to other dispatched ops, so it works on any backend that has those primitives. It is used as a fallback when no backend-specific impl is registered.

nkipy.core.ops.get_backend() → str

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

Set the current source location in the trace context.

Parameters:

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