"""
Matrix operations API for OpenFHE-Numpy.
This module provides NumPy-compatible matrix operations that can be performed on
encrypted data using homomorphic encryption. Functions follow NumPy naming
conventions and similar signatures where possible.
All functions use the tensor_function_api decorator to handle different tensor types
and dispatch to the appropriate backend implementation.
"""
from typing import Any, Optional, Union
from numpy.typing import ArrayLike
from .dispatch import tensor_function_api
# ===========================
# Element-wise Operations
# ===========================
[docs]
def add(a: ArrayLike, b: ArrayLike) -> ArrayLike:
"""
Add two tensors or a tensor and a scalar.
Parameters
----------
a : ArrayLike
First operand.
b : ArrayLike
Second operand.
Returns
-------
out : ArrayLike
Element-wise sum of `a` and `b`.
See Also
--------
numpy.add : Corresponding NumPy function.
Examples
--------
>>> add([1, 2], [3, 4])
array([4, 6])
"""
return _add_dispatch(a, b)
@tensor_function_api("add", binary=True)
def _add_dispatch(a: ArrayLike, b: ArrayLike) -> ArrayLike:
pass
[docs]
@tensor_function_api("subtract", binary=True)
def subtract(a: ArrayLike, b: ArrayLike) -> ArrayLike:
"""
Subtract two tensors or a tensor and a scalar.
Parameters
----------
a : ArrayLike
First operand.
b : ArrayLike
Second operand.
Returns
-------
out : ArrayLike
Element-wise difference of `a` and `b`.
See Also
--------
numpy.subtract : Corresponding NumPy function.
Examples
--------
>>> subtract([5, 7], [2, 4])
array([3, 3])
"""
pass
[docs]
@tensor_function_api("multiply", binary=True)
def multiply(a: ArrayLike, b: ArrayLike) -> ArrayLike:
"""
Element-wise multiply two tensors or a tensor and a scalar.
Parameters
----------
a : ArrayLike
First operand.
b : ArrayLike or scalar
Second operand.
Returns
-------
out : ArrayLike
Element-wise product of `a` and `b`.
See Also
--------
numpy.multiply : Corresponding NumPy function.
Examples
--------
>>> multiply([1, 2], [3, 4])
array([3, 8])
"""
pass
[docs]
@tensor_function_api("power", binary=True)
def power(a: ArrayLike, exponent: int) -> ArrayLike:
"""
For each element of the tensor, it raises that element to the given power.
Note
----
This only supports integer exponents due to homomorphic-encryption constraints.
Parameters
----------
a : ArrayLike
Base tensor.
exponent : int
Power to raise each array element to.
Returns
-------
out : ArrayLike
Element-wise `a` raised to `exponent`.
See Also
--------
numpy.power : Corresponding element-wise power function.
numpy.linalg.matrix_power : Repeated matrix multiplication for square matrices.
Examples
--------
>>> power([1, 2, 3], 2)
array([1, 4, 9])
"""
pass
# ===========================
# Matrix Operations
# ===========================
[docs]
@tensor_function_api("dot", binary=True)
def dot(a: ArrayLike, b: ArrayLike) -> ArrayLike:
"""
Compute the dot (inner) product of two tensors.
1-D vectors: inner product
2-D matrices: matrix multiplication
Parameters
----------
a : ArrayLike
First operand.
b : ArrayLike
Second operand.
Returns
-------
out : ArrayLike
Dot product of `a` and `b`.
See Also
--------
numpy.dot : Corresponding NumPy function.
Examples
--------
# 1-D vectors: inner product
>>> dot([1, 2], [3, 4])
11
# 2-D matrices: matrix multiplication
>>> import numpy as np
>>> A = np.array([[1, 2], [3, 4]])
>>> B = np.array([[5, 6], [7, 8]])
>>> dot(A, B)
array([[19, 22],
[43, 50]])
"""
pass
[docs]
@tensor_function_api("matmul", binary=True)
def matmul(a: ArrayLike, b: ArrayLike) -> ArrayLike:
"""
Matrix multiply two tensors.
Parameters
----------
a : ArrayLike
First operand.
b : ArrayLike
Second operand.
Returns
-------
out : ArrayLike
Matrix product of `a` and `b`.
See Also
--------
numpy.matmul : Corresponding NumPy function.
Examples
--------
>>> import numpy as np
>>> matmul(np.array([[1, 2], [3, 4]]), np.array([[5, 6], [7, 8]]))
array([[19, 22],
[43, 50]])
"""
pass
[docs]
@tensor_function_api("transpose", binary=False)
def transpose(a: ArrayLike) -> ArrayLike:
"""
Transpose a tensor.
Parameters
----------
a : ArrayLike
Input tensor.
Returns
-------
out : ArrayLike
Transposed tensor.
See Also
--------
numpy.transpose : Corresponding NumPy function.
Examples
--------
>>> import numpy as np
>>> transpose(np.array([[1, 2], [3, 4]]))
array([[1, 3],
[2, 4]])
"""
pass
# ===========================
# Reduction Operations
# ===========================
[docs]
@tensor_function_api("cumsum", binary=False)
def cumsum(a: ArrayLike, axis: int = 0, keepdims: bool = False) -> ArrayLike:
"""
Compute the cumulative sum of tensor elements along an axis.
Parameters
----------
a : ArrayLike
Input tensor.
axis : int, optional
Axis along which to compute the sum. Default is 0.
keepdims : bool, optional
If True, retains reduced dimensions. Default is False.
Returns
-------
out : ArrayLike
Cumulative sum of `a`.
See Also
--------
numpy.cumsum : Corresponding NumPy function.
Examples
--------
>>> import numpy as np
>>> cumsum(np.array([[1, 2], [3, 4]]), axis=1)
array([[1, 3],
[3, 7]])
"""
pass
[docs]
@tensor_function_api("cumreduce", binary=False)
def cumreduce(a: ArrayLike, axis: int = 0, keepdims: bool = False) -> ArrayLike:
"""
Compute the cumulative reduction (e.g., product) of tensor elements.
Parameters
----------
a : ArrayLike
Input tensor.
axis : int, optional
Axis along which to compute the reduction. Default is 0.
keepdims : bool, optional
If True, retains reduced dimensions. Default is False.
Returns
-------
out : ArrayLike
Cumulative reduction of `a`.
See Also
--------
numpy.cumprod : Similar operation for product.
Examples
--------
>>> import numpy as np
>>> cumreduce(np.array([[1, 2, 3], [4, 5, 6]]), axis=0)
array([[1, 2, 3],
[-3, -3, -3]])
"""
pass
[docs]
@tensor_function_api("sum", binary=False)
def sum(a: ArrayLike, axis: Optional[int] = 0, keepdims: bool = False) -> ArrayLike:
"""
Sum of tensor elements over a given axis.
Parameters
----------
a : ArrayLike
Input tensor.
axis : int, optional
Axis along which to compute the sum. Default is 0.
0: sum over rows
1: sum over cols
keepdims : bool, optional
If True, retains reduced dimensions. Default is False.
Returns
-------
out : ArrayLike
Sum of `a` elements.
See Also
--------
numpy.sum : Corresponding NumPy function.
Examples
--------
>>> import numpy as np
>>> a = np.array([[1, 2],
[3, 4]])
>>> sum(a)
10
>>> sum(a, axis=0)
array([4, 6])
>>> sum(a, axis=1)
array([3, 7])
"""
pass
[docs]
@tensor_function_api("mean", binary=False)
def mean(a: ArrayLike, axis: Optional[int] = 0, keepdims: bool = False) -> ArrayLike:
"""
Compute the arithmetic mean along the specified axis.
Returns the average of the array elements. The average is taken over
the flattened array by default, otherwise over the specified axis.
Parameters
----------
a : ArrayLike
Input tensor.
axis : int, optional
Axis along which to compute the mean. Default is 0.
keepdims : bool, optional
If True, retains reduced dimensions. Default is False.
Returns
-------
out : ArrayLike
Mean of `a` elements.
See Also
--------
numpy.mean : Corresponding NumPy function.
Examples
--------
>>> import numpy as np
>>> a = np.array([[1, 2], [3, 4]])
>>> mean(a)
2.5
>>> mean(a, axis=0)
array([2., 3.])
>>> mean(a, axis=1)
array([1.5, 3.5])
"""
pass
# ===========================
# Planned Future Functionality
# ===========================
# Array Creation Functions:
# def zeros(shape, crypto_context, key):
# """Create an encrypted array of zeros."""
# pass
# def ones(shape, crypto_context, key):
# """Create an encrypted array of ones."""
# pass
# def eye(n, crypto_context, key):
# """Create an encrypted identity matrix."""
# pass
# Broadcasting Support:
# def _get_broadcast_shape(self, other):
# """Calculate broadcast shape between tensors."""
# # Implementation...
# def _broadcast_tensor(self, target_shape):
# """Broadcast this tensor to target shape."""
# # Implementation...
# Array Manipulation Methods:
# def reshape(self, new_shape):
# """Reshape tensor to new dimensions."""
# # Implementation...
# def concat(tensors, axis=0):
# """Concatenate tensors along specified axis."""
# # Implementation...