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.

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`
`cos`	hlo	`np.cos`
`exp`	hlo	`np.exp`
`floor`	hlo	`np.floor`
`invert`	hlo	`np.invert`
`log`	hlo	`np.log`
`logical_not`	hlo	`np.logical_not`
`negative`	hlo	`np.negative`
`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`

Linear Algebra Operations#

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

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`

Collective Operations#

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

Convolution Operations#

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

Neural Network Operations#

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

Creation Operations#

Operation	Backend(s)	NumPy Equivalent
`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`

Reduction Operations#

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

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`
`repeat`	hlo	`np.repeat`
`reshape`	hlo	`np.reshape`
`split`	hlo	`np.split`
`transpose`	hlo	`np.transpose`

Indexing Operations#

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

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.set_backend(backend: str | None, context: TraceContext | None = None) → None[source]#

Set the current backend and optional trace context.

Called by tracing infrastructure when a kernel is being traced.

Parameters:

backend – Backend name (‘hlo’, ‘cpu’, or future backends), or None to clear.
context – Trace context for tracing backends (required for ‘hlo’).

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

Get the current backend name.

Returns:: Current backend name.
Raises:: RuntimeError – If no backend is set.

nkipy.core.ops.set_context(context: TraceContext | None) → None[source]#

Set the current trace context.

This is typically called together with set_backend, but can be used separately if needed.

Parameters:: context – Trace context or None to clear.

nkipy.core.ops.get_context() → TraceContext | None[source]#

Get the current trace context.

Returns:: Current trace context, or None for eager backends (like ‘cpu’).

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

Set the current source location in the trace context.

This is a convenience function that sets the source location in the current context if available.

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

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

Get the current source location from the trace context.

This is a convenience function that retrieves the source location from the current context if available.

Returns:: Tuple of (filename, line_number) or None if no context or location set.

Operations Reference

Contents