Add Op

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.backend.hlo import get_hlo_context
from nkipy.core.ops._registry import Op

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

@my_op.impl('hlo')
def _my_op_hlo(arg1, arg2, ...):
    """HLO tracing implementation."""
    ctx = get_hlo_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_hlo_context() from nkipy.core.backend.hlo)
• 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

Converted and distributed by TomeVault — claim your Tome and manage your conversions.