Add a New NKIPy Operation

Add the operation $0 to the $1 category file.

Steps

1. Identify the correct category file

Operations live in nkipy/src/nkipy/core/ops/<category>.py. Choose from:

2. Create the Op instance and register implementations

In the chosen category file, follow this pattern:

from nkipy.core.ops._registry import Op, get_context

# Create the dispatcher
my_op = Op('my_op')

@my_op.impl('hlo')
def _my_op_hlo(arg1, arg2, ...):
    """HLO tracing implementation."""
    ctx = get_context()
    # Build HLO operation using ctx
    # Return NKIPyTensorRef
    ...

@my_op.impl('cpu')
def _my_op_cpu(arg1, arg2, ...):
    """CPU eager implementation using NumPy."""
    import numpy as np
    # Return numpy array
    ...

Key rules:

• •The hlo impl builds HLO IR nodes via the trace context (get_context())
• •The cpu impl uses pure NumPy and returns numpy arrays
• •Both implementations must accept the same signature and produce equivalent results
• •Both backends are required — every op needs hlo and cpu

3. Export from __init__.py

Add the new op to nkipy/src/nkipy/core/ops/__init__.py:

• •Import it from the category module
• •Add it to the __all__ list if one exists
• •Follow the existing grouping/ordering in the imports

4. Wire into tensor API (if applicable)

If the op should be callable as a method on tensors (e.g., tensor.sum()), or via np.* dispatch:

• •Check nkipy/src/nkipy/core/tensor.py for TensorArithmeticMixin
• •Check nkipy/src/nkipy/core/tensor_apis.py for high-level API wrappers
• •Check nkipy/src/nkipy/core/_numpy_dispatch.py for NumPy protocol integration

5. Add tests

Create or extend tests in tests/unit/. Follow existing patterns:

• •Test both CPU and HLO backends
• •Test edge cases (empty tensors, broadcasting, dtype promotion)
• •Run with: uv run pytest tests/unit/test_tensor_api.py -k "test_my_op" -v

6. Verify

uv run ruff check nkipy/src/nkipy/core/ops/
uv run pytest tests/unit/ -k "$0" -v