mirror of https://github.com/Nonannet/copapy.git
Nicolas Kruse
GitHub
commit
9f60a22d3d
|
|
@ -39,6 +39,16 @@ jobs:
|
|||
name: stencil-object-files
|
||||
path: src/copapy/obj
|
||||
|
||||
- uses: actions/download-artifact@v4
|
||||
with:
|
||||
name: musl-object-files
|
||||
path: /tmp/musl-object-files
|
||||
|
||||
- name: Add musl copyright notice to license file
|
||||
run: |
|
||||
echo "\n\nMUSL COPYRIGHT NOTICE:" >> LICENSE
|
||||
cat /tmp/musl-object-files/COPYRIGHT >> LICENSE
|
||||
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@v5
|
||||
with:
|
||||
|
|
|
|||
|
|
@ -28,7 +28,7 @@ jobs:
|
|||
- uses: actions/upload-artifact@v4
|
||||
with:
|
||||
name: musl-object-files
|
||||
path: /object_files/musl_objects_*.*o
|
||||
path: /object_files/*
|
||||
|
||||
build-package-test:
|
||||
needs: [build_stencils]
|
||||
|
|
@ -47,6 +47,16 @@ jobs:
|
|||
name: stencil-object-files
|
||||
path: src/copapy/obj
|
||||
|
||||
- uses: actions/download-artifact@v4
|
||||
with:
|
||||
name: musl-object-files
|
||||
path: /tmp/musl-object-files
|
||||
|
||||
- name: Add musl copyright notice to license file
|
||||
run: |
|
||||
echo "\n\nMUSL COPYRIGHT NOTICE:" >> LICENSE
|
||||
cat /tmp/musl-object-files/COPYRIGHT >> LICENSE
|
||||
|
||||
- name: Set up Python ${{ matrix.python-version }}
|
||||
uses: actions/setup-python@v5
|
||||
with:
|
||||
|
|
@ -289,7 +299,8 @@ jobs:
|
|||
set -v
|
||||
mkdir -p release
|
||||
cp tmp/stencil-object-files/* release/
|
||||
cp tmp/musl-object-files/* release/
|
||||
cp tmp/musl-object-files/*.o release/
|
||||
cp tmp/musl-object-files/COPYRIGHT release/MUSL-COPYRIGHT.txt
|
||||
cp tmp/runner-linux-x86_64/coparun release/
|
||||
cp tmp/runner-linux-arm64/coparun release/coparun-aarch64
|
||||
cp tmp/runner-linux-armv6/coparun release/coparun-armv6
|
||||
|
|
|
|||
25
README.md
25
README.md
|
|
@ -13,7 +13,7 @@ The main features can be summarized as:
|
|||
- Memory and type safety with a minimal set of runtime errors
|
||||
- Deterministic execution
|
||||
- Automatic differentiation for efficient realtime optimization (reverse-mode)
|
||||
- Optimized machine code for x86_64, AArch64 and ARMv7
|
||||
- Optimized machine code for x86_64, ARMv6, ARMv7 and AArch64
|
||||
- Highly portable to new architectures
|
||||
- Small Python package with minimal dependencies and no cross-compile toolchain required
|
||||
|
||||
|
|
@ -38,10 +38,29 @@ Despite missing SIMD-optimization, benchmark performance shows promising numbers
|
|||
|
||||
|
||||
|
||||
For the benchmark (`tests/benchmark.py`) the timing of 30000 iterations for calculating the therm `sum((v1 + i) @ v2 for i in range(10))` where measured on an Ryzen 5 3400G. Where the vectors `v1` and `v2` both have a lengths of `v_size` which was varied according to the chart from 10 to 600. For the NumPy case the "i in range(10)" loop was vectorized like this: `np.sum((v1 + i) @ v2)` with i being here a `NDArray` with a dimension of `[10, 1]`. The number of calculated scalar operations is the same for both contenders. Obviously copapy profits from less overheat by calling a single function from python per iteration, where the NumPy variant requires 3. Interestingly there is no indication visible in the chart that for increasing `v_size` the calling overhead for NumPy will be compensated by using faster SIMD instructions.
|
||||
For the benchmark (`tests/benchmark.py`) the timing of 30000 iterations for calculating the therm `sum((v1 + i) @ v2 for i in range(10))` where measured on an Ryzen 5 3400G. Where the vectors `v1` and `v2` both have a lengths of `v_size` which was varied according to the chart from 10 to 600. For the NumPy case the "i in range(10)" loop was vectorized like this: `np.sum((v1 + i) @ v2)` with i being here a `NDArray` with a dimension of `[10, 1]`. The number of calculated scalar operations is the same for both contenders. Obviously copapy profits from less overheat by calling a single function from python per iteration, where the NumPy variant requires 3. Interestingly there is no indication visible in the chart that for increasing `v_size` the calling overhead for NumPy will be compensated by using faster SIMD instructions. It is to note that in this benchmark the copapy case does not move any data between python and the compiled code.
|
||||
|
||||
Furthermore for many applications copypy will benefit by reducing the actual number of operations significantly compared to a NumPy implementation, by precompute constant values know at compile time and benefiting from sparcity. Multiplying by zero (e.g. in a diagonal matrix) eliminate a hole branch in the computation graph. Operations without effect, like multiplications by 1 oder additions with zero gets eliminated at compile time.
|
||||
|
||||
For Testing and using Copapy to speed up computations in conventional Python programs there is also the `@cp.jit` decorator available, to compile functions on first use and cache the compiled version for later calls:
|
||||
|
||||
```python
|
||||
import copapy as cp
|
||||
|
||||
@cp.jit
|
||||
def calculation(x: float, y: float) -> float:
|
||||
return sum(x ** 2 + y ** 2 + i for i in range(10))
|
||||
|
||||
# Compile and run:
|
||||
result1 = calculation(2.5, 1.2)
|
||||
|
||||
# Run cached compiled version:
|
||||
result2 = calculation(3.1, 4.7)
|
||||
```
|
||||
|
||||
It is to note that `cp.jit` is not optimized very much at the moment concerning transfer data between Python and the compiled code back and forth.
|
||||
|
||||
|
||||
## Install
|
||||
|
||||
To install Copapy, you can use pip. Precompiled wheels are available for Linux (x86_64, AArch64, ARMv7), Windows (x86_64) and macOS (x86_64, AArch64):
|
||||
|
|
@ -234,4 +253,4 @@ This project is licensed under the MIT license - see the [LICENSE](LICENSE) file
|
|||
|
||||
[^2]: The compiler must support tail-call optimization (TCO). Currently, GCC is supported. Porting to a new architecture requires implementing a subset of relocation types used by that architecture.
|
||||
|
||||
[^3]: Supported architectures: x86_64, AArch64, ARMv7 (non-Thumb). ARMv6/7-M (Thumb) support is in development. Code for x86 32-bit exists but has unresolved issues and a low priority.
|
||||
[^3]: Supported architectures: x86_64, AArch64, ARMv6 and 7 (non-Thumb). ARMv6/7-M (Thumb) support is in development. Code for x86 32-bit exists but has unresolved issues and a low priority.
|
||||
|
|
|
|||
|
|
@ -34,7 +34,7 @@ def build_asm_code_dict(asm_glob_pattern: str) -> dict[str, str]:
|
|||
# Example usage:
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(description="Generate stencils documentation from C and assembly code")
|
||||
parser.add_argument('--input', default='tools/make_example.py', help='Path to input C file')
|
||||
parser.add_argument('--input', default='tools/make_example.py', help='Path to example script')
|
||||
parser.add_argument('--asm-pattern', default='build/tmp/runner-linux-*/example.asm', help='Glob pattern for assembly files')
|
||||
parser.add_argument('--output', default='docs/build/compiled_example.md', help='Output markdown file path')
|
||||
|
||||
|
|
|
|||
|
|
@ -134,7 +134,8 @@ if __name__ == "__main__":
|
|||
md_code: str = ''
|
||||
|
||||
for function_name, code in functions.items():
|
||||
md_code += get_stencil_section(function_name)
|
||||
if 'get_42' not in function_name and not function_name.startswith('cast_'):
|
||||
md_code += get_stencil_section(function_name)
|
||||
|
||||
with open(args.output, 'wt') as f:
|
||||
f.write(md_code)
|
||||
|
|
|
|||
|
|
@ -36,9 +36,11 @@ Example usage:
|
|||
from ._target import Target, jit
|
||||
from ._basic_types import NumLike, value, generic_sdb, iif
|
||||
from ._vectors import vector, distance, scalar_projection, angle_between, rotate_vector, vector_projection
|
||||
from ._matrices import matrix, identity, zeros, ones, diagonal, eye
|
||||
from ._tensors import tensor, zeros, ones, arange, eye, identity, diagonal
|
||||
from ._math import sqrt, abs, sign, sin, cos, tan, asin, acos, atan, atan2, log, exp, pow, get_42, clamp, min, max, relu
|
||||
from ._autograd import grad
|
||||
from ._tensors import tensor as matrix
|
||||
|
||||
|
||||
__all__ = [
|
||||
"Target",
|
||||
|
|
@ -47,11 +49,13 @@ __all__ = [
|
|||
"generic_sdb",
|
||||
"iif",
|
||||
"vector",
|
||||
"tensor",
|
||||
"matrix",
|
||||
"identity",
|
||||
"zeros",
|
||||
"ones",
|
||||
"diagonal",
|
||||
"arange",
|
||||
"sqrt",
|
||||
"abs",
|
||||
"sin",
|
||||
|
|
|
|||
|
|
@ -1,4 +1,4 @@
|
|||
from . import value, vector, matrix
|
||||
from . import value, vector, tensor
|
||||
import copapy.backend as cpb
|
||||
from typing import Any, Sequence, overload
|
||||
import copapy as cp
|
||||
|
|
@ -10,10 +10,10 @@ def grad(x: Any, y: value[Any]) -> unifloat: ...
|
|||
@overload
|
||||
def grad(x: Any, y: vector[Any]) -> vector[float]: ...
|
||||
@overload
|
||||
def grad(x: Any, y: Sequence[value[Any]]) -> list[unifloat]: ...
|
||||
def grad(x: Any, y: tensor[Any]) -> tensor[float]: ...
|
||||
@overload
|
||||
def grad(x: Any, y: matrix[Any]) -> matrix[float]: ...
|
||||
def grad(x: Any, y: value[Any] | Sequence[value[Any]] | vector[Any] | matrix[Any]) -> Any:
|
||||
def grad(x: Any, y: Sequence[value[Any]]) -> list[unifloat]: ...
|
||||
def grad(x: Any, y: value[Any] | Sequence[value[Any]] | vector[Any] | tensor[Any]) -> Any:
|
||||
"""Returns the partial derivative dx/dy where x needs to be a scalar
|
||||
and y might be a scalar, a list of scalars, a vector or matrix. It
|
||||
uses automatic differentiation in reverse-mode.
|
||||
|
|
@ -29,11 +29,11 @@ def grad(x: Any, y: value[Any] | Sequence[value[Any]] | vector[Any] | matrix[Any
|
|||
|
||||
if isinstance(y, value):
|
||||
y_set = {y}
|
||||
if isinstance(y, matrix):
|
||||
y_set = {v for row in y for v in row}
|
||||
if isinstance(y, tensor):
|
||||
y_set = {v.get_scalar(0) for v in y.flatten()}
|
||||
else:
|
||||
assert isinstance(y, Sequence) or isinstance(y, vector)
|
||||
y_set = {v for v in y}
|
||||
y_set = set(y)
|
||||
|
||||
edges = cpb.get_all_dag_edges_between([x.net.source], (v.net.source for v in y_set if isinstance(v, value)))
|
||||
ordered_ops = cpb.stable_toposort(edges)
|
||||
|
|
@ -121,7 +121,7 @@ def grad(x: Any, y: value[Any] | Sequence[value[Any]] | vector[Any] | matrix[Any
|
|||
if isinstance(y, value):
|
||||
return grad_dict[y.net]
|
||||
if isinstance(y, vector):
|
||||
return vector(grad_dict[yi.net] if isinstance(yi, value) else 0.0 for yi in y)
|
||||
if isinstance(y, matrix):
|
||||
return matrix((grad_dict[yi.net] if isinstance(yi, value) else 0.0 for yi in row) for row in y)
|
||||
return vector(grad_dict[yi.net] if isinstance(yi, value) else 0.0 for yi in y.values)
|
||||
if isinstance(y, tensor):
|
||||
return tensor([grad_dict[yi.net] if isinstance(yi, value) else 0.0 for yi in y.values], y.shape)
|
||||
return [grad_dict[yi.net] for yi in y]
|
||||
|
|
|
|||
|
|
@ -1,5 +1,5 @@
|
|||
import pkgutil
|
||||
from typing import Any, Sequence, TypeVar, overload, TypeAlias, Generic, cast
|
||||
from typing import Any, Sequence, TypeVar, overload, TypeAlias, Generic, cast, Callable
|
||||
from ._stencils import stencil_database, detect_process_arch
|
||||
import copapy as cp
|
||||
from ._helper_types import TNum
|
||||
|
|
@ -75,7 +75,7 @@ class Net:
|
|||
|
||||
def __hash__(self) -> int:
|
||||
return self.source.node_hash
|
||||
|
||||
|
||||
def __eq__(self, other: object) -> bool:
|
||||
return isinstance(other, Net) and self.source == other.source
|
||||
|
||||
|
|
@ -396,7 +396,7 @@ class Op(Node):
|
|||
return True
|
||||
if not isinstance(other, Op):
|
||||
return NotImplemented
|
||||
|
||||
|
||||
# Traverse graph for both notes. Return false on first difference.
|
||||
# A false inequality result in seldom cases is ok, whereas a false
|
||||
# equality result leads to wrong computation results.
|
||||
|
|
@ -427,10 +427,19 @@ class Op(Node):
|
|||
return False
|
||||
seen.add(key)
|
||||
return True
|
||||
|
||||
|
||||
def __hash__(self) -> int:
|
||||
return self.node_hash
|
||||
|
||||
# Interface for vector and tensor types
|
||||
class ArrayType(Generic[TNum]):
|
||||
def __init__(self, shape: tuple[int, ...]) -> None:
|
||||
self.shape = shape
|
||||
self.values: tuple[TNum | value[TNum], ...] = ()
|
||||
|
||||
def map(self, func: Callable[[TNum | value[TNum]], Any]) -> 'ArrayType[Any]':
|
||||
return self
|
||||
|
||||
|
||||
def value_from_number(val: Any) -> value[Any]:
|
||||
# Create anonymous constant that can be removed during optimization
|
||||
|
|
@ -454,7 +463,7 @@ def iif(expression: float | int | value[Any], true_result: TNum | value[TNum], f
|
|||
def iif(expression: Any, true_result: Any, false_result: Any) -> Any:
|
||||
"""Inline if-else operation. Returns true_result if expression is non-zero,
|
||||
else returns false_result.
|
||||
|
||||
|
||||
Arguments:
|
||||
expression: The condition to evaluate.
|
||||
true_result: The result if expression is non-zero.
|
||||
|
|
|
|||
|
|
@ -1,359 +0,0 @@
|
|||
from . import value
|
||||
from ._vectors import vector
|
||||
from ._mixed import mixed_sum
|
||||
from typing import TypeVar, Iterable, Any, overload, TypeAlias, Callable, Iterator, Generic
|
||||
from ._helper_types import TNum
|
||||
|
||||
MatNumLike: TypeAlias = 'matrix[int] | matrix[float] | value[int] | value[float] | int | float'
|
||||
MatIntLike: TypeAlias = 'matrix[int] | value[int] | int'
|
||||
MatFloatLike: TypeAlias = 'matrix[float] | value[float] | float'
|
||||
U = TypeVar("U", int, float)
|
||||
|
||||
|
||||
class matrix(Generic[TNum]):
|
||||
"""Mathematical matrix class supporting basic operations and interactions with values.
|
||||
"""
|
||||
def __init__(self, values: Iterable[Iterable[TNum | value[TNum]]] | vector[TNum]):
|
||||
"""Create a matrix with given values.
|
||||
|
||||
Arguments:
|
||||
values: iterable of iterable of constant values
|
||||
"""
|
||||
if isinstance(values, vector):
|
||||
rows = [values.values]
|
||||
else:
|
||||
rows = [tuple(row) for row in values]
|
||||
|
||||
if rows:
|
||||
row_len = len(rows[0])
|
||||
assert all(len(row) == row_len for row in rows), "All rows must have the same length"
|
||||
self.values: tuple[tuple[value[TNum] | TNum, ...], ...] = tuple(rows)
|
||||
self.rows = len(self.values)
|
||||
self.cols = len(self.values[0]) if self.values else 0
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"matrix({self.values})"
|
||||
|
||||
def __len__(self) -> int:
|
||||
"""Return the number of rows in the matrix."""
|
||||
return self.rows
|
||||
|
||||
@overload
|
||||
def __getitem__(self, key: int) -> vector[TNum]: ...
|
||||
@overload
|
||||
def __getitem__(self, key: tuple[int, int]) -> value[TNum] | TNum: ...
|
||||
def __getitem__(self, key: int | tuple[int, int]) -> Any:
|
||||
"""Get a row as a vector or a specific element.
|
||||
Arguments:
|
||||
key: row index or (row, col) tuple
|
||||
|
||||
Returns:
|
||||
vector if row index is given, else the element at (row, col)
|
||||
"""
|
||||
if isinstance(key, tuple):
|
||||
assert len(key) == 2
|
||||
return self.values[key[0]][key[1]]
|
||||
else:
|
||||
return vector(self.values[key])
|
||||
|
||||
def __iter__(self) -> Iterator[tuple[value[TNum] | TNum, ...]]:
|
||||
return iter(self.values)
|
||||
|
||||
def __neg__(self) -> 'matrix[TNum]':
|
||||
return matrix((-a for a in row) for row in self.values)
|
||||
|
||||
@overload
|
||||
def __add__(self: 'matrix[int]', other: MatFloatLike) -> 'matrix[float]': ...
|
||||
@overload
|
||||
def __add__(self: 'matrix[int]', other: MatIntLike) -> 'matrix[int]': ...
|
||||
@overload
|
||||
def __add__(self: 'matrix[float]', other: MatNumLike) -> 'matrix[float]': ...
|
||||
@overload
|
||||
def __add__(self, other: MatNumLike) -> 'matrix[int] | matrix[float]': ...
|
||||
def __add__(self, other: MatNumLike) -> Any:
|
||||
if isinstance(other, matrix):
|
||||
assert self.rows == other.rows and self.cols == other.cols, \
|
||||
"Matrices must have the same dimensions"
|
||||
return matrix(
|
||||
tuple(a + b for a, b in zip(row1, row2))
|
||||
for row1, row2 in zip(self.values, other.values)
|
||||
)
|
||||
if isinstance(other, value):
|
||||
return matrix(
|
||||
tuple(a + other for a in row)
|
||||
for row in self.values
|
||||
)
|
||||
o = value(other) # Make sure a single constant is allocated
|
||||
return matrix(
|
||||
tuple(a + o if isinstance(a, value) else a + other for a in row)
|
||||
for row in self.values
|
||||
)
|
||||
|
||||
@overload
|
||||
def __radd__(self: 'matrix[float]', other: MatNumLike) -> 'matrix[float]': ...
|
||||
@overload
|
||||
def __radd__(self: 'matrix[int]', other: value[int] | int) -> 'matrix[int]': ...
|
||||
def __radd__(self, other: Any) -> Any:
|
||||
return self + other
|
||||
|
||||
@overload
|
||||
def __sub__(self: 'matrix[int]', other: MatFloatLike) -> 'matrix[float]': ...
|
||||
@overload
|
||||
def __sub__(self: 'matrix[int]', other: MatIntLike) -> 'matrix[int]': ...
|
||||
@overload
|
||||
def __sub__(self: 'matrix[float]', other: MatNumLike) -> 'matrix[float]': ...
|
||||
@overload
|
||||
def __sub__(self, other: MatNumLike) -> 'matrix[int] | matrix[float]': ...
|
||||
def __sub__(self, other: MatNumLike) -> Any:
|
||||
if isinstance(other, matrix):
|
||||
assert self.rows == other.rows and self.cols == other.cols, \
|
||||
"Matrices must have the same dimensions"
|
||||
return matrix(
|
||||
tuple(a - b for a, b in zip(row1, row2))
|
||||
for row1, row2 in zip(self.values, other.values)
|
||||
)
|
||||
if isinstance(other, value):
|
||||
return matrix(
|
||||
tuple(a - other for a in row)
|
||||
for row in self.values
|
||||
)
|
||||
o = value(other) # Make sure a single constant is allocated
|
||||
return matrix(
|
||||
tuple(a - o if isinstance(a, value) else a - other for a in row)
|
||||
for row in self.values
|
||||
)
|
||||
|
||||
@overload
|
||||
def __rsub__(self: 'matrix[float]', other: MatNumLike) -> 'matrix[float]': ...
|
||||
@overload
|
||||
def __rsub__(self: 'matrix[int]', other: value[int] | int) -> 'matrix[int]': ...
|
||||
def __rsub__(self, other: MatNumLike) -> Any:
|
||||
if isinstance(other, matrix):
|
||||
assert self.rows == other.rows and self.cols == other.cols, \
|
||||
"Matrices must have the same dimensions"
|
||||
return matrix(
|
||||
tuple(b - a for a, b in zip(row1, row2))
|
||||
for row1, row2 in zip(self.values, other.values)
|
||||
)
|
||||
if isinstance(other, value):
|
||||
return matrix(
|
||||
tuple(other - a for a in row)
|
||||
for row in self.values
|
||||
)
|
||||
o = value(other) # Make sure a single constant is allocated
|
||||
return matrix(
|
||||
tuple(o - a if isinstance(a, value) else other - a for a in row)
|
||||
for row in self.values
|
||||
)
|
||||
|
||||
@overload
|
||||
def __mul__(self: 'matrix[int]', other: MatFloatLike) -> 'matrix[float]': ...
|
||||
@overload
|
||||
def __mul__(self: 'matrix[int]', other: MatIntLike) -> 'matrix[int]': ...
|
||||
@overload
|
||||
def __mul__(self: 'matrix[float]', other: MatNumLike) -> 'matrix[float]': ...
|
||||
@overload
|
||||
def __mul__(self, other: MatNumLike) -> 'matrix[int] | matrix[float]': ...
|
||||
def __mul__(self, other: MatNumLike) -> Any:
|
||||
"""Element-wise multiplication"""
|
||||
if isinstance(other, matrix):
|
||||
assert self.rows == other.rows and self.cols == other.cols, \
|
||||
"Matrices must have the same dimensions"
|
||||
return matrix(
|
||||
tuple(a * b for a, b in zip(row1, row2))
|
||||
for row1, row2 in zip(self.values, other.values)
|
||||
)
|
||||
if isinstance(other, value):
|
||||
return matrix(
|
||||
tuple(a * other for a in row)
|
||||
for row in self.values
|
||||
)
|
||||
o = value(other) # Make sure a single constant is allocated
|
||||
return matrix(
|
||||
tuple(a * o if isinstance(a, value) else a * other for a in row)
|
||||
for row in self.values
|
||||
)
|
||||
|
||||
@overload
|
||||
def __rmul__(self: 'matrix[float]', other: MatNumLike) -> 'matrix[float]': ...
|
||||
@overload
|
||||
def __rmul__(self: 'matrix[int]', other: value[int] | int) -> 'matrix[int]': ...
|
||||
def __rmul__(self, other: MatNumLike) -> Any:
|
||||
return self * other
|
||||
|
||||
def __truediv__(self, other: MatNumLike) -> 'matrix[float]':
|
||||
"""Element-wise division"""
|
||||
if isinstance(other, matrix):
|
||||
assert self.rows == other.rows and self.cols == other.cols, \
|
||||
"Matrices must have the same dimensions"
|
||||
return matrix(
|
||||
tuple(a / b for a, b in zip(row1, row2))
|
||||
for row1, row2 in zip(self.values, other.values)
|
||||
)
|
||||
if isinstance(other, value):
|
||||
return matrix(
|
||||
tuple(a / other for a in row)
|
||||
for row in self.values
|
||||
)
|
||||
o = value(other) # Make sure a single constant is allocated
|
||||
return matrix(
|
||||
tuple(a / o if isinstance(a, value) else a / other for a in row)
|
||||
for row in self.values
|
||||
)
|
||||
|
||||
def __rtruediv__(self, other: MatNumLike) -> 'matrix[float]':
|
||||
if isinstance(other, matrix):
|
||||
assert self.rows == other.rows and self.cols == other.cols, \
|
||||
"Matrices must have the same dimensions"
|
||||
return matrix(
|
||||
tuple(b / a for a, b in zip(row1, row2))
|
||||
for row1, row2 in zip(self.values, other.values)
|
||||
)
|
||||
if isinstance(other, value):
|
||||
return matrix(
|
||||
tuple(other / a for a in row)
|
||||
for row in self.values
|
||||
)
|
||||
o = value(other) # Make sure a single constant is allocated
|
||||
return matrix(
|
||||
tuple(o / a if isinstance(a, value) else other / a for a in row)
|
||||
for row in self.values
|
||||
)
|
||||
|
||||
@overload
|
||||
def __matmul__(self: 'matrix[TNum]', other: 'vector[TNum]') -> 'vector[TNum]': ...
|
||||
@overload
|
||||
def __matmul__(self: 'matrix[TNum]', other: 'matrix[TNum]') -> 'matrix[TNum]': ...
|
||||
def __matmul__(self: 'matrix[TNum]', other: 'matrix[TNum] | vector[TNum]') -> 'matrix[TNum] | vector[TNum]':
|
||||
"""Matrix multiplication using @ operator"""
|
||||
if isinstance(other, vector):
|
||||
assert self.cols == len(other.values), \
|
||||
f"Matrix columns ({self.cols}) must match vector length ({len(other.values)})"
|
||||
vec_result = (mixed_sum(a * b for a, b in zip(row, other.values)) for row in self.values)
|
||||
return vector(vec_result)
|
||||
else:
|
||||
assert isinstance(other, matrix), "Cannot multiply matrix with {type(other)}"
|
||||
assert self.cols == other.rows, \
|
||||
f"Matrix columns ({self.cols}) must match other matrix rows ({other.rows})"
|
||||
result: list[list[TNum | value[TNum]]] = []
|
||||
for row in self.values:
|
||||
new_row: list[TNum | value[TNum]] = []
|
||||
for col_idx in range(other.cols):
|
||||
col = tuple(other.values[i][col_idx] for i in range(other.rows))
|
||||
element = sum(a * b for a, b in zip(row, col))
|
||||
new_row.append(element)
|
||||
result.append(new_row)
|
||||
return matrix(result)
|
||||
|
||||
def transpose(self) -> 'matrix[TNum]':
|
||||
"""Return the transpose of the matrix."""
|
||||
if not self.values:
|
||||
return matrix([])
|
||||
return matrix(
|
||||
tuple(self.values[i][j] for i in range(self.rows))
|
||||
for j in range(self.cols)
|
||||
)
|
||||
|
||||
@property
|
||||
def shape(self) -> tuple[int, int]:
|
||||
"""Return the shape of the matrix as (rows, cols)."""
|
||||
return (self.rows, self.cols)
|
||||
|
||||
@property
|
||||
def T(self) -> 'matrix[TNum]':
|
||||
return self.transpose()
|
||||
|
||||
def row(self, index: int) -> vector[TNum]:
|
||||
"""Get a row as a vector."""
|
||||
assert 0 <= index < self.rows, f"Row index {index} out of bounds"
|
||||
return vector(self.values[index])
|
||||
|
||||
def col(self, index: int) -> vector[TNum]:
|
||||
"""Get a column as a vector."""
|
||||
assert 0 <= index < self.cols, f"Column index {index} out of bounds"
|
||||
return vector(self.values[i][index] for i in range(self.rows))
|
||||
|
||||
@overload
|
||||
def trace(self: 'matrix[TNum]') -> TNum | value[TNum]: ...
|
||||
@overload
|
||||
def trace(self: 'matrix[int]') -> int | value[int]: ...
|
||||
@overload
|
||||
def trace(self: 'matrix[float]') -> float | value[float]: ...
|
||||
def trace(self) -> Any:
|
||||
"""Calculate the trace (sum of diagonal elements)."""
|
||||
assert self.rows == self.cols, "Trace is only defined for square matrices"
|
||||
return mixed_sum(self.values[i][i] for i in range(self.rows))
|
||||
|
||||
@overload
|
||||
def sum(self: 'matrix[TNum]') -> TNum | value[TNum]: ...
|
||||
@overload
|
||||
def sum(self: 'matrix[int]') -> int | value[int]: ...
|
||||
@overload
|
||||
def sum(self: 'matrix[float]') -> float | value[float]: ...
|
||||
def sum(self) -> Any:
|
||||
"""Calculate the sum of all elements."""
|
||||
return mixed_sum(a for row in self.values for a in row)
|
||||
|
||||
def map(self, func: Callable[[Any], value[U] | U]) -> 'matrix[U]':
|
||||
"""Applies a function to each element of the matrix and returns a new matrix."""
|
||||
return matrix(
|
||||
tuple(func(a) for a in row)
|
||||
for row in self.values
|
||||
)
|
||||
|
||||
def homogenize(self) -> 'matrix[TNum]':
|
||||
"""Convert all elements to copapy values if any element is a copapy value."""
|
||||
if any(isinstance(val, value) for row in self.values for val in row):
|
||||
return matrix(
|
||||
tuple(value(val) if not isinstance(val, value) else val for val in row)
|
||||
for row in self.values
|
||||
)
|
||||
else:
|
||||
return self
|
||||
|
||||
|
||||
def identity(size: int) -> matrix[int]:
|
||||
"""Create an identity matrix of given size."""
|
||||
return matrix(
|
||||
tuple(1 if i == j else 0 for j in range(size))
|
||||
for i in range(size)
|
||||
)
|
||||
|
||||
|
||||
def zeros(rows: int, cols: int) -> matrix[int]:
|
||||
"""Create a zero matrix of given dimensions."""
|
||||
return matrix(
|
||||
tuple(0 for _ in range(cols))
|
||||
for _ in range(rows)
|
||||
)
|
||||
|
||||
|
||||
def ones(rows: int, cols: int) -> matrix[int]:
|
||||
"""Create a matrix of ones with given dimensions."""
|
||||
return matrix(
|
||||
tuple(1 for _ in range(cols))
|
||||
for _ in range(rows)
|
||||
)
|
||||
|
||||
|
||||
def eye(rows: int, cols: int | None = None) -> matrix[int]:
|
||||
"""Create a matrix with ones on the diagonal and zeros elsewhere."""
|
||||
cols = cols if cols else rows
|
||||
return matrix(
|
||||
tuple(1 if i == j else 0 for j in range(cols))
|
||||
for i in range(rows)
|
||||
)
|
||||
|
||||
|
||||
@overload
|
||||
def diagonal(vec: 'vector[int]') -> matrix[int]: ...
|
||||
@overload
|
||||
def diagonal(vec: 'vector[float]') -> matrix[float]: ...
|
||||
def diagonal(vec: vector[Any]) -> matrix[Any]:
|
||||
"""Create a diagonal matrix from a vector."""
|
||||
size = len(vec)
|
||||
|
||||
return matrix(
|
||||
tuple(vec[i] if i == j else 0 for j in range(size))
|
||||
for i in range(size)
|
||||
)
|
||||
|
|
@ -2,8 +2,7 @@ from typing import Iterable, overload, TypeVar, Any, Callable, TypeAlias
|
|||
from . import _binwrite as binw
|
||||
from coparun_module import coparun, read_data_mem, create_target, clear_target
|
||||
import struct
|
||||
from ._basic_types import stencil_db_from_package
|
||||
from ._basic_types import value, Net, Node, Write, NumLike
|
||||
from ._basic_types import value, Net, Node, Write, NumLike, ArrayType, stencil_db_from_package
|
||||
from ._compiler import compile_to_dag
|
||||
|
||||
T = TypeVar("T", int, float)
|
||||
|
|
@ -66,7 +65,7 @@ class Target():
|
|||
def __del__(self) -> None:
|
||||
clear_target(self._context)
|
||||
|
||||
def compile(self, *values: int | float | value[Any] | Iterable[int | float | value[Any]]) -> None:
|
||||
def compile(self, *values: NumLike | value[T] | ArrayType[T] | Iterable[T | value[T]]) -> None:
|
||||
"""Compiles the code to compute the given values.
|
||||
|
||||
Arguments:
|
||||
|
|
@ -74,13 +73,16 @@ class Target():
|
|||
"""
|
||||
nodes: list[Node] = []
|
||||
for input in values:
|
||||
if isinstance(input, Iterable):
|
||||
if isinstance(input, ArrayType):
|
||||
for v in input.values:
|
||||
if isinstance(v, value):
|
||||
nodes.append(Write(v))
|
||||
elif isinstance(input, Iterable):
|
||||
for v in input:
|
||||
if isinstance(v, value):
|
||||
nodes.append(Write(v))
|
||||
else:
|
||||
if isinstance(input, value):
|
||||
nodes.append(Write(input))
|
||||
elif isinstance(input, value):
|
||||
nodes.append(Write(input))
|
||||
|
||||
dw, self._values = compile_to_dag(nodes, self.sdb)
|
||||
dw.write_com(binw.Command.END_COM)
|
||||
|
|
@ -100,7 +102,9 @@ class Target():
|
|||
def read_value(self, variables: NumLike) -> float | int | bool: ...
|
||||
@overload
|
||||
def read_value(self, variables: Iterable[T | value[T]]) -> list[T]: ...
|
||||
def read_value(self, variables: NumLike | value[T] | Iterable[T | value[T]]) -> Any:
|
||||
@overload
|
||||
def read_value(self, variables: ArrayType[T]) -> ArrayType[T]: ...
|
||||
def read_value(self, variables: NumLike | value[T] | ArrayType[T] | Iterable[T | value[T]]) -> Any:
|
||||
"""Reads the numeric value of a copapy type.
|
||||
|
||||
Arguments:
|
||||
|
|
@ -109,6 +113,9 @@ class Target():
|
|||
Returns:
|
||||
Numeric value or values
|
||||
"""
|
||||
if isinstance(variables, ArrayType):
|
||||
return variables.map(lambda v: self.read_value(v))
|
||||
|
||||
if isinstance(variables, Iterable):
|
||||
return [self.read_value(ni) if isinstance(ni, value) else ni for ni in variables]
|
||||
|
||||
|
|
@ -142,7 +149,7 @@ class Target():
|
|||
return val
|
||||
else:
|
||||
raise ValueError(f"Unsupported value type: {var_type}")
|
||||
|
||||
|
||||
def write_value(self, variables: value[Any] | Iterable[value[Any]], data: int | float | Iterable[int | float]) -> None:
|
||||
"""Write to a copapy value on the target.
|
||||
|
||||
|
|
@ -155,7 +162,7 @@ class Target():
|
|||
for ni, vi in zip(variables, data):
|
||||
self.write_value(ni, vi)
|
||||
return
|
||||
|
||||
|
||||
assert not isinstance(data, Iterable), "If net is not iterable, value must not be iterable"
|
||||
|
||||
assert isinstance(variables, value), "Argument must be a copapy value"
|
||||
|
|
@ -174,7 +181,7 @@ class Target():
|
|||
dw.write_value(int(data), lengths)
|
||||
else:
|
||||
raise ValueError(f"Unsupported value type: {var_type}")
|
||||
|
||||
|
||||
dw.write_com(binw.Command.END_COM)
|
||||
assert coparun(self._context, dw.get_data()) > 0
|
||||
|
||||
|
|
|
|||
|
|
@ -0,0 +1,938 @@
|
|||
from copapy._basic_types import NumLike, ArrayType
|
||||
from . import value
|
||||
from ._vectors import vector
|
||||
from ._mixed import mixed_sum
|
||||
from typing import TypeVar, Any, overload, TypeAlias, Callable, Iterator, Sequence
|
||||
from ._helper_types import TNum
|
||||
|
||||
TensorNumLike: TypeAlias = 'tensor[Any] | vector[Any] | value[Any] | int | float | bool'
|
||||
TensorIntLike: TypeAlias = 'tensor[int] | value[int] | int'
|
||||
TensorFloatLike: TypeAlias = 'tensor[float] | value[float] | float'
|
||||
TensorSequence: TypeAlias = 'Sequence[TNum | value[TNum]] | Sequence[Sequence[TNum | value[TNum]]] | Sequence[Sequence[Sequence[TNum | value[TNum]]]]'
|
||||
U = TypeVar("U", int, float)
|
||||
|
||||
|
||||
class tensor(ArrayType[TNum]):
|
||||
"""Generalized n-dimensional tensor class supporting numpy-style operations.
|
||||
|
||||
A tensor can have any number of dimensions and supports element-wise operations,
|
||||
reshaping, transposition, and various reduction operations.
|
||||
"""
|
||||
|
||||
def __init__(self, values: 'TNum | value[TNum] | vector[TNum] | tensor[TNum] | TensorSequence[TNum]', shape: Sequence[int] | None = None):
|
||||
"""Create a tensor with given values.
|
||||
|
||||
Arguments:
|
||||
values: Nested iterables of constant values or copapy values.
|
||||
Can be a scalar, 1D iterable (vector),
|
||||
or n-dimensional nested structure.
|
||||
"""
|
||||
if shape:
|
||||
self.shape: tuple[int, ...] = tuple(shape)
|
||||
assert (isinstance(values, Sequence) and
|
||||
any(isinstance(v, (value, int, float)) for v in values)), \
|
||||
"Values must be a sequence of values if shape is provided"
|
||||
self.values: tuple[TNum | value[TNum], ...] = tuple(v for v in values if not isinstance(v, Sequence))
|
||||
self.ndim: int = len(shape)
|
||||
elif isinstance(values, (int, float)):
|
||||
# Scalar case: 0-dimensional tensor
|
||||
self.shape = ()
|
||||
self.values = (values,)
|
||||
self.ndim = 0
|
||||
elif isinstance(values, value):
|
||||
# Scalar value case
|
||||
self.shape = ()
|
||||
self.values = (values,)
|
||||
self.ndim = 0
|
||||
elif isinstance(values, vector):
|
||||
# 1D case from vector
|
||||
self.shape = (len(values),)
|
||||
self.values = values.values
|
||||
self.ndim = 1
|
||||
elif isinstance(values, tensor):
|
||||
# Copy constructor
|
||||
self.shape = values.shape
|
||||
self.values = values.values
|
||||
self.ndim = values.ndim
|
||||
else:
|
||||
# General n-dimensional case
|
||||
self.values, self.shape = self._infer_shape_and_flatten(values)
|
||||
self.ndim = len(self.shape)
|
||||
|
||||
@staticmethod
|
||||
def _infer_shape_and_flatten(values: Sequence[Any]) -> tuple[tuple[Any, ...], tuple[int, ...]]:
|
||||
"""Infer the shape of a nested iterable and validate consistency."""
|
||||
def get_shape(val: int | float | value[Any] | Sequence[Any]) -> list[int]:
|
||||
if isinstance(val, int | float):
|
||||
return []
|
||||
if isinstance(val, value):
|
||||
return []
|
||||
else:
|
||||
if not val:
|
||||
return [0]
|
||||
sub_shape = get_shape(val[0])
|
||||
if any(get_shape(item) != sub_shape for item in val[1:]):
|
||||
raise ValueError("All elements must have consistent shape")
|
||||
return [len(val)] + sub_shape
|
||||
return []
|
||||
|
||||
shape = tuple(get_shape(values))
|
||||
if not shape:
|
||||
# Scalar
|
||||
return (values,), ()
|
||||
|
||||
# Flatten nested structure
|
||||
def flatten_recursive(val: Any) -> list[Any]:
|
||||
if isinstance(val, int | float | value):
|
||||
return [val]
|
||||
else:
|
||||
result: list[value[Any]] = []
|
||||
for item in val:
|
||||
if isinstance(item, int | float | value | Sequence):
|
||||
result.extend(flatten_recursive(item))
|
||||
return result
|
||||
|
||||
flattened = flatten_recursive(values)
|
||||
return tuple(flattened), shape
|
||||
|
||||
def _get_flat_index(self, indices: Sequence[int]) -> int:
|
||||
"""Convert multi-dimensional indices to flat index."""
|
||||
if len(indices) != len(self.shape):
|
||||
raise IndexError(f"Expected {len(self.shape)} indices, got {len(indices)}")
|
||||
|
||||
flat_idx = 0
|
||||
stride = 1
|
||||
for i in range(len(self.shape) - 1, -1, -1):
|
||||
if not (0 <= indices[i] < self.shape[i]):
|
||||
raise IndexError(f"Index {indices[i]} out of bounds for dimension {i} with size {self.shape[i]}")
|
||||
flat_idx += indices[i] * stride
|
||||
stride *= self.shape[i]
|
||||
return flat_idx
|
||||
|
||||
def _get_indices_from_flat(self, flat_idx: int) -> tuple[int, ...]:
|
||||
"""Convert flat index to multi-dimensional indices."""
|
||||
indices: list[int] = []
|
||||
for dim_size in reversed(self.shape):
|
||||
indices.append(flat_idx % dim_size)
|
||||
flat_idx //= dim_size
|
||||
return tuple(reversed(indices))
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"tensor(shape={self.shape}, values={self.values if self.ndim == 0 else '...'})"
|
||||
|
||||
def __len__(self) -> int:
|
||||
"""Return the size of the first dimension."""
|
||||
if self.ndim == 0:
|
||||
raise TypeError("len() of a 0-d tensor")
|
||||
return self.shape[0]
|
||||
|
||||
def get_scalar(self: 'tensor[TNum]', *key: int) -> TNum | value[TNum]:
|
||||
"""Get a single scalar value from the tensor given multi-dimensional indices."""
|
||||
assert len(key) == self.ndim, f"Expected {self.ndim} indices, got {len(key)}"
|
||||
flat_idx = self._get_flat_index(key)
|
||||
return self.values[flat_idx]
|
||||
|
||||
def __getitem__(self, key: int | slice | Sequence[int | slice]) -> 'tensor[TNum]':
|
||||
"""Get a sub-tensor or element.
|
||||
|
||||
Arguments:
|
||||
key: Integer index (returns tensor of rank n-1),
|
||||
slice object (returns tensor with same rank),
|
||||
tuple of indices/slices (returns sub-tensor or element),
|
||||
or tuple of indices (returns single element).
|
||||
|
||||
Returns:
|
||||
Sub-tensor or element value.
|
||||
"""
|
||||
if self.ndim == 0:
|
||||
raise TypeError("Cannot index a 0-d tensor")
|
||||
|
||||
# Handle single slice
|
||||
if isinstance(key, slice):
|
||||
return self._handle_slice((key,))
|
||||
|
||||
# Handle tuple of indices/slices
|
||||
if isinstance(key, Sequence):
|
||||
return self._handle_slice(key)
|
||||
|
||||
# Handle single integer index
|
||||
assert isinstance(key, int), f"indices must be integers, slices, or tuples thereof, not {type(key)}"
|
||||
# Return a sub-tensor of rank n-1
|
||||
if not (-self.shape[0] <= key < self.shape[0]):
|
||||
raise IndexError(f"Index {key} out of bounds for dimension 0 with size {self.shape[0]}")
|
||||
|
||||
if key < 0:
|
||||
key += self.shape[0]
|
||||
|
||||
# Calculate which elements belong to this slice
|
||||
sub_shape = self.shape[1:]
|
||||
sub_size = 1
|
||||
for s in sub_shape:
|
||||
sub_size *= s
|
||||
|
||||
start_idx = key * sub_size
|
||||
end_idx = start_idx + sub_size
|
||||
|
||||
sub_values = self.values[start_idx:end_idx]
|
||||
|
||||
if not sub_shape:
|
||||
#assert False, (sub_shape, len(sub_shape), sub_values[0])
|
||||
return tensor(sub_values[0])
|
||||
|
||||
return tensor(sub_values, sub_shape)
|
||||
|
||||
def _handle_slice(self, keys: Sequence[int | slice]) -> 'tensor[TNum]':
|
||||
"""Handle slicing operations on the tensor."""
|
||||
# Process all keys and identify ranges for each dimension
|
||||
ranges: list[range] = []
|
||||
|
||||
for i, key in enumerate(keys):
|
||||
if i >= self.ndim:
|
||||
raise IndexError(f"Too many indices for tensor of rank {self.ndim}")
|
||||
|
||||
dim_size = self.shape[i]
|
||||
|
||||
if isinstance(key, int):
|
||||
if not (-dim_size <= key < dim_size):
|
||||
raise IndexError(f"Index {key} out of bounds for dimension {i} with size {dim_size}")
|
||||
if key < 0:
|
||||
key += dim_size
|
||||
ranges.append(range(key, key + 1))
|
||||
else:
|
||||
assert isinstance(key, slice), f"indices must be integers or slices, not {type(key)}"
|
||||
start, stop, step = key.indices(dim_size)
|
||||
ranges.append(range(start, stop, step))
|
||||
|
||||
# Handle remaining dimensions (full ranges)
|
||||
for i in range(len(keys), self.ndim):
|
||||
ranges.append(range(self.shape[i]))
|
||||
|
||||
# Collect elements matching the ranges
|
||||
selected_values: list[TNum | value[TNum]] = []
|
||||
new_shape: list[int] = []
|
||||
|
||||
# Calculate new shape (only include dimensions that weren't single integers)
|
||||
for i, key in enumerate(keys):
|
||||
if not isinstance(key, int):
|
||||
new_shape.append(len(ranges[i]))
|
||||
|
||||
# Add remaining dimensions
|
||||
for i in range(len(keys), self.ndim):
|
||||
new_shape.append(self.shape[i])
|
||||
|
||||
# Iterate through all combinations of indices in the ranges
|
||||
def iterate_ranges(range_list: list[range], current_indices: list[int]) -> None:
|
||||
if len(current_indices) == len(range_list):
|
||||
# Compute flat index
|
||||
flat_idx = 0
|
||||
stride = 1
|
||||
for i in range(len(self.shape) - 1, -1, -1):
|
||||
flat_idx += current_indices[i] * stride
|
||||
stride *= self.shape[i]
|
||||
selected_values.append(self.values[flat_idx])
|
||||
else:
|
||||
dim = len(current_indices)
|
||||
for idx in range_list[dim]:
|
||||
iterate_ranges(range_list, current_indices + [idx])
|
||||
|
||||
iterate_ranges(ranges, [])
|
||||
|
||||
# Return based on result shape
|
||||
if not new_shape:
|
||||
# Single element (all were integers)
|
||||
return tensor(selected_values[0])
|
||||
|
||||
return tensor(tuple(selected_values), tuple(new_shape))
|
||||
|
||||
def __iter__(self) -> Iterator['tensor[TNum]']:
|
||||
"""Iterate over the first dimension."""
|
||||
if self.ndim == 0:
|
||||
raise TypeError("Cannot iterate over a 0-d tensor")
|
||||
|
||||
for i in range(self.shape[0]):
|
||||
yield self[i]
|
||||
|
||||
def __neg__(self) -> 'tensor[TNum]':
|
||||
"""Negate all elements."""
|
||||
negated_values: tuple[Any, ...] = tuple(-v for v in self.values)
|
||||
return tensor(negated_values, self.shape)
|
||||
|
||||
@overload
|
||||
def __add__(self: 'tensor[int]', other: TensorFloatLike) -> 'tensor[float]': ...
|
||||
@overload
|
||||
def __add__(self: 'tensor[int]', other: TensorIntLike) -> 'tensor[int]': ...
|
||||
@overload
|
||||
def __add__(self: 'tensor[float]', other: TensorNumLike) -> 'tensor[float]': ...
|
||||
@overload
|
||||
def __add__(self, other: TensorNumLike) -> 'tensor[int] | tensor[float]': ...
|
||||
def __add__(self, other: TensorNumLike) -> Any:
|
||||
"""Element-wise addition."""
|
||||
return self._binary_op(other, lambda a, b: a + b)
|
||||
|
||||
@overload
|
||||
def __radd__(self: 'tensor[float]', other: TensorNumLike) -> 'tensor[float]': ...
|
||||
@overload
|
||||
def __radd__(self: 'tensor[int]', other: value[int] | int) -> 'tensor[int]': ...
|
||||
def __radd__(self, other: Any) -> Any:
|
||||
return self + other
|
||||
|
||||
@overload
|
||||
def __sub__(self: 'tensor[int]', other: TensorFloatLike) -> 'tensor[float]': ...
|
||||
@overload
|
||||
def __sub__(self: 'tensor[int]', other: TensorIntLike) -> 'tensor[int]': ...
|
||||
@overload
|
||||
def __sub__(self: 'tensor[float]', other: TensorNumLike) -> 'tensor[float]': ...
|
||||
@overload
|
||||
def __sub__(self, other: TensorNumLike) -> 'tensor[int] | tensor[float]': ...
|
||||
def __sub__(self, other: TensorNumLike) -> Any:
|
||||
"""Element-wise subtraction."""
|
||||
return self._binary_op(other, lambda a, b: a - b, commutative=False)
|
||||
|
||||
@overload
|
||||
def __rsub__(self: 'tensor[float]', other: TensorNumLike) -> 'tensor[float]': ...
|
||||
@overload
|
||||
def __rsub__(self: 'tensor[int]', other: value[int] | int) -> 'tensor[int]': ...
|
||||
def __rsub__(self, other: TensorNumLike) -> Any:
|
||||
return self._binary_op(other, lambda a, b: b - a, commutative=False, reversed=True)
|
||||
|
||||
@overload
|
||||
def __mul__(self: 'tensor[int]', other: TensorFloatLike) -> 'tensor[float]': ...
|
||||
@overload
|
||||
def __mul__(self: 'tensor[int]', other: TensorIntLike) -> 'tensor[int]': ...
|
||||
@overload
|
||||
def __mul__(self: 'tensor[float]', other: TensorNumLike) -> 'tensor[float]': ...
|
||||
@overload
|
||||
def __mul__(self, other: TensorNumLike) -> 'tensor[int] | tensor[float]': ...
|
||||
def __mul__(self, other: TensorNumLike) -> Any:
|
||||
"""Element-wise multiplication."""
|
||||
return self._binary_op(other, lambda a, b: a * b)
|
||||
|
||||
@overload
|
||||
def __rmul__(self: 'tensor[float]', other: TensorNumLike) -> 'tensor[float]': ...
|
||||
@overload
|
||||
def __rmul__(self: 'tensor[int]', other: value[int] | int) -> 'tensor[int]': ...
|
||||
def __rmul__(self, other: TensorNumLike) -> Any:
|
||||
return self * other
|
||||
|
||||
def __truediv__(self, other: TensorNumLike) -> 'tensor[float]':
|
||||
"""Element-wise division."""
|
||||
return self._binary_op(other, lambda a, b: a / b, commutative=False)
|
||||
|
||||
def __rtruediv__(self, other: TensorNumLike) -> 'tensor[float]':
|
||||
"""Element-wise right division."""
|
||||
return self._binary_op(other, lambda a, b: b / a, commutative=False, reversed=True)
|
||||
|
||||
@overload
|
||||
def __pow__(self: 'tensor[int]', other: TensorFloatLike) -> 'tensor[float]': ...
|
||||
@overload
|
||||
def __pow__(self: 'tensor[int]', other: TensorIntLike) -> 'tensor[int]': ...
|
||||
@overload
|
||||
def __pow__(self: 'tensor[float]', other: TensorNumLike) -> 'tensor[float]': ...
|
||||
@overload
|
||||
def __pow__(self, other: TensorNumLike) -> 'tensor[int] | tensor[float]': ...
|
||||
def __pow__(self, other: TensorNumLike) -> Any:
|
||||
"""Element-wise power."""
|
||||
return self._binary_op(other, lambda a, b: a ** b, commutative=False)
|
||||
|
||||
@overload
|
||||
def __rpow__(self: 'tensor[float]', other: TensorNumLike) -> 'tensor[float]': ...
|
||||
@overload
|
||||
def __rpow__(self: 'tensor[int]', other: value[int] | int) -> 'tensor[int]': ...
|
||||
def __rpow__(self, other: TensorNumLike) -> Any:
|
||||
return self._binary_op(other, lambda a, b: b ** a, commutative=False, reversed=True)
|
||||
|
||||
def __gt__(self, other: TensorNumLike) -> 'tensor[int]':
|
||||
"""Element-wise greater than."""
|
||||
return self._binary_op(other, lambda a, b: a > b, commutative=False)
|
||||
|
||||
def __lt__(self, other: TensorNumLike) -> 'tensor[int]':
|
||||
"""Element-wise less than."""
|
||||
return self._binary_op(other, lambda a, b: a < b, commutative=False)
|
||||
|
||||
def __ge__(self, other: TensorNumLike) -> 'tensor[int]':
|
||||
"""Element-wise greater than or equal."""
|
||||
return self._binary_op(other, lambda a, b: a >= b, commutative=False)
|
||||
|
||||
def __le__(self, other: TensorNumLike) -> 'tensor[int]':
|
||||
"""Element-wise less than or equal."""
|
||||
return self._binary_op(other, lambda a, b: a <= b, commutative=False)
|
||||
|
||||
def __eq__(self, other: TensorNumLike) -> 'tensor[int]': # type: ignore
|
||||
"""Element-wise equality."""
|
||||
return self._binary_op(other, lambda a, b: a == b)
|
||||
|
||||
def __ne__(self, other: TensorNumLike) -> 'tensor[int]': # type: ignore
|
||||
"""Element-wise inequality."""
|
||||
return self._binary_op(other, lambda a, b: a != b)
|
||||
|
||||
def _binary_op(self, other: TensorNumLike, op: Callable[[Any, Any], 'tensor[TNum]'],
|
||||
commutative: bool = True, reversed: bool = False) -> 'tensor[Any]':
|
||||
"""Perform binary operation with broadcasting support.
|
||||
"""
|
||||
seen_consts: dict[NumLike, NumLike] = {}
|
||||
|
||||
def call_op(a: TNum | value[TNum], b: NumLike) -> Any:
|
||||
if isinstance(b, value) or not isinstance(a, value):
|
||||
b_trans = b
|
||||
else:
|
||||
if b in seen_consts:
|
||||
b_trans = seen_consts[b]
|
||||
else:
|
||||
b_trans = value(b)
|
||||
seen_consts[b] = b_trans
|
||||
if reversed:
|
||||
return op(b_trans, a)
|
||||
else:
|
||||
return op(a, b_trans)
|
||||
|
||||
if isinstance(other, Sequence | vector):
|
||||
other_tensor: tensor[Any] = tensor(other)
|
||||
return self._binary_op(other_tensor, op, commutative, reversed)
|
||||
|
||||
elif isinstance(other, tensor):
|
||||
self_shape = self.shape
|
||||
other_shape = other.shape
|
||||
|
||||
# Check if shapes are identical
|
||||
if self_shape == other_shape:
|
||||
result_vals = tuple(call_op(a, b) for a, b in zip(self.values, other.values))
|
||||
return tensor(result_vals, self_shape)
|
||||
|
||||
# Broadcast shapes using numpy-style broadcasting rules
|
||||
result_shape = self._broadcast_shapes(self_shape, other_shape)
|
||||
|
||||
# Expand both tensors to the broadcast shape
|
||||
self_expanded = self._expand_to_shape(result_shape)
|
||||
other_expanded = other._expand_to_shape(result_shape)
|
||||
|
||||
# Apply operation element-wise
|
||||
result_vals = tuple(call_op(a, b) for a, b in zip(self_expanded.values, other_expanded.values))
|
||||
return tensor(result_vals, result_shape)
|
||||
|
||||
else:
|
||||
# Broadcast scalar
|
||||
result_vals = tuple(call_op(v, other) for v in self.values)
|
||||
return tensor(result_vals, self.shape)
|
||||
|
||||
def _broadcast_shapes(self, shape1: tuple[int, ...], shape2: tuple[int, ...]) -> tuple[int, ...]:
|
||||
"""Compute the broadcast shape of two shapes following numpy rules.
|
||||
|
||||
Rules:
|
||||
- Dimensions are compared from right to left
|
||||
- Dimensions must either be equal or one must be 1
|
||||
- Missing dimensions are treated as 1
|
||||
"""
|
||||
# Align from the right
|
||||
max_ndim = max(len(shape1), len(shape2))
|
||||
|
||||
# Pad with 1s on the left
|
||||
padded_shape1 = (1,) * (max_ndim - len(shape1)) + shape1
|
||||
padded_shape2 = (1,) * (max_ndim - len(shape2)) + shape2
|
||||
|
||||
result_shape: list[int] = []
|
||||
for dim1, dim2 in zip(padded_shape1, padded_shape2):
|
||||
if dim1 == dim2:
|
||||
result_shape.append(dim1)
|
||||
elif dim1 == 1:
|
||||
result_shape.append(dim2)
|
||||
elif dim2 == 1:
|
||||
result_shape.append(dim1)
|
||||
else:
|
||||
raise ValueError(f"Incompatible shapes for broadcasting: {shape1} vs {shape2}")
|
||||
|
||||
return tuple(result_shape)
|
||||
|
||||
def _expand_to_shape(self, target_shape: tuple[int, ...]) -> 'tensor[TNum]':
|
||||
"""Expand tensor to target shape using broadcasting (repeating dimensions of size 1).
|
||||
"""
|
||||
if self.shape == target_shape:
|
||||
return self
|
||||
|
||||
# Pad self.shape with 1s on the left to match target_shape length
|
||||
padded_self_shape = (1,) * (len(target_shape) - len(self.shape)) + self.shape
|
||||
|
||||
# Validate broadcasting is possible
|
||||
for s, t in zip(padded_self_shape, target_shape):
|
||||
if s != t and s != 1:
|
||||
raise ValueError(f"Cannot broadcast shape {self.shape} to {target_shape}")
|
||||
|
||||
# Expand step by step from left to right
|
||||
current_tensor = self
|
||||
current_shape = self.shape
|
||||
|
||||
# Add missing dimensions on the left
|
||||
if len(self.shape) < len(target_shape):
|
||||
diff = len(target_shape) - len(self.shape)
|
||||
for _ in range(diff):
|
||||
# Reshape to add dimension of size 1 on the left
|
||||
current_tensor = current_tensor.reshape(1, *current_tensor.shape)
|
||||
current_shape = (1,) + current_shape
|
||||
|
||||
# Expand each dimension that is 1 to the target size
|
||||
for i, (curr_dim, target_dim) in enumerate(zip(current_shape, target_shape)):
|
||||
if curr_dim == 1 and target_dim != 1:
|
||||
current_tensor = current_tensor._repeat_along_axis(i, target_dim)
|
||||
current_shape = current_tensor.shape
|
||||
|
||||
return current_tensor
|
||||
|
||||
def _repeat_along_axis(self, axis: int, repetitions: int) -> 'tensor[TNum]':
|
||||
"""Repeat tensor along a specific axis.
|
||||
"""
|
||||
if self.shape[axis] != 1:
|
||||
raise ValueError(f"Can only repeat dimensions of size 1, got {self.shape[axis]}")
|
||||
|
||||
# Create list of indices to select (repeat the single index)
|
||||
new_shape = list(self.shape)
|
||||
new_shape[axis] = repetitions
|
||||
new_values: list[TNum | value[TNum]] = []
|
||||
|
||||
# Iterate through all positions in the new tensor
|
||||
def iterate_and_repeat(current_indices: list[int], depth: int) -> None:
|
||||
if depth == len(new_shape):
|
||||
# Get the source index (with 0 for the repeated dimension)
|
||||
source_indices = list(current_indices)
|
||||
source_indices[axis] = 0
|
||||
source_idx = self._get_flat_index(tuple(source_indices))
|
||||
new_values.append(self.values[source_idx])
|
||||
else:
|
||||
for i in range(new_shape[depth]):
|
||||
iterate_and_repeat(current_indices + [i], depth + 1)
|
||||
|
||||
iterate_and_repeat([], 0)
|
||||
return tensor(tuple(new_values), tuple(new_shape))
|
||||
|
||||
def reshape(self, *new_shape: int) -> 'tensor[TNum]':
|
||||
"""Reshape the tensor to a new shape.
|
||||
|
||||
Arguments:
|
||||
*new_shape: New shape dimensions. Use -1 for one dimension to infer automatically.
|
||||
|
||||
Returns:
|
||||
A new tensor with the specified shape.
|
||||
"""
|
||||
shape_arg: tuple[int, ...] | int = new_shape if len(new_shape) != 1 else new_shape[0]
|
||||
if isinstance(shape_arg, int):
|
||||
new_shape = (shape_arg,)
|
||||
else:
|
||||
new_shape = shape_arg
|
||||
|
||||
# Handle -1 in shape (automatic dimension inference)
|
||||
neg_one_count = sum(1 for d in new_shape if d == -1)
|
||||
if neg_one_count > 1:
|
||||
raise ValueError("Only one dimension can be -1")
|
||||
|
||||
if neg_one_count == 1:
|
||||
known_size = 1
|
||||
for dim in new_shape:
|
||||
if dim != -1:
|
||||
known_size *= dim
|
||||
|
||||
if self.size() % known_size != 0:
|
||||
raise ValueError(f"Cannot infer dimension from size {self.size()} with shape {new_shape}")
|
||||
|
||||
inferred_dim = self.size() // known_size
|
||||
new_shape = tuple(inferred_dim if d == -1 else d for d in new_shape)
|
||||
|
||||
total_size = 1
|
||||
for dim in new_shape:
|
||||
total_size *= dim
|
||||
|
||||
if total_size != self.size():
|
||||
raise ValueError(f"Cannot reshape tensor of size {self.size()} into shape {new_shape}")
|
||||
|
||||
return tensor(self.values, new_shape)
|
||||
|
||||
@overload
|
||||
def trace(self: 'tensor[TNum]') -> TNum | value[TNum]: ...
|
||||
@overload
|
||||
def trace(self: 'tensor[int]') -> int | value[int]: ...
|
||||
@overload
|
||||
def trace(self: 'tensor[float]') -> float | value[float]: ...
|
||||
def trace(self) -> Any:
|
||||
"""Calculate the trace (sum of diagonal elements)."""
|
||||
assert self.ndim == 2 and self.shape[0] == self.shape[1], "Trace is only defined for square matrices"
|
||||
return mixed_sum(self.get_scalar(i, i) for i in range(self.shape[0]))
|
||||
|
||||
def transpose(self, *axes: int) -> 'tensor[TNum]':
|
||||
"""Transpose the tensor.
|
||||
|
||||
Arguments:
|
||||
*axes: Permutation of axes. If not provided, reverses all axes.
|
||||
|
||||
Returns:
|
||||
A transposed tensor.
|
||||
"""
|
||||
if not axes:
|
||||
axes = tuple(range(self.ndim - 1, -1, -1))
|
||||
|
||||
if len(axes) != self.ndim:
|
||||
raise ValueError("axes don't match tensor")
|
||||
|
||||
if any(not (0 <= ax < self.ndim) for ax in axes):
|
||||
raise ValueError(f"Invalid axes for tensor of rank {self.ndim}")
|
||||
|
||||
new_shape = tuple(self.shape[ax] for ax in axes)
|
||||
new_values: list[Any] = [None] * len(self.values)
|
||||
|
||||
for old_idx in range(len(self.values)):
|
||||
old_indices = self._get_indices_from_flat(old_idx)
|
||||
new_indices = tuple(old_indices[ax] for ax in axes)
|
||||
|
||||
new_flat_idx = 0
|
||||
stride = 1
|
||||
for i in range(len(new_shape) - 1, -1, -1):
|
||||
new_flat_idx += new_indices[i] * stride
|
||||
stride *= new_shape[i]
|
||||
|
||||
new_values[new_flat_idx] = self.values[old_idx]
|
||||
|
||||
return tensor(new_values, new_shape)
|
||||
|
||||
def flatten(self) -> 'tensor[TNum]':
|
||||
"""Flatten the tensor to 1D.
|
||||
|
||||
Returns:
|
||||
A flattened 1D tensor.
|
||||
"""
|
||||
return self.reshape(-1)
|
||||
|
||||
def size(self) -> int:
|
||||
"""Return total number of elements."""
|
||||
size = 1
|
||||
for dim in self.shape:
|
||||
size *= dim
|
||||
return size
|
||||
|
||||
def matmul(self, other: 'tensor[TNum] | vector[TNum]') -> 'TNum | value[TNum] | tensor[TNum]':
|
||||
"""Matrix multiplication (@ operator).
|
||||
|
||||
Arguments:
|
||||
other: Another tensor to multiply with.
|
||||
|
||||
Returns:
|
||||
Result of matrix multiplication.
|
||||
|
||||
Raises:
|
||||
ValueError: If shapes are incompatible for matrix multiplication.
|
||||
"""
|
||||
if self.ndim < 1 or other.ndim < 1:
|
||||
raise ValueError("matmul requires tensors with at least 1 dimension")
|
||||
|
||||
# For 1D x 1D: dot product (returns scalar)
|
||||
if self.ndim == 1 and other.ndim == 1:
|
||||
if self.shape[0] != other.shape[0]:
|
||||
raise ValueError(f"Shape mismatch: ({self.shape[0]},) @ ({other.shape[0]},)")
|
||||
result = mixed_sum(a * b for a, b in zip(self.values, other.values))
|
||||
return result
|
||||
|
||||
# For 2D x 2D: standard matrix multiplication
|
||||
if self.ndim == 2 and other.ndim == 2 and isinstance(other, tensor):
|
||||
if self.shape[1] != other.shape[0]:
|
||||
raise ValueError(f"Shape mismatch: {self.shape} @ {other.shape}")
|
||||
|
||||
result_values: list[Any] = []
|
||||
for i in range(self.shape[0]):
|
||||
for j in range(other.shape[1]):
|
||||
dot_sum = sum(self.get_scalar(i, k) * other.get_scalar(k, j) for k in range(self.shape[1]))
|
||||
result_values.append(dot_sum)
|
||||
|
||||
return tensor(tuple(result_values), (self.shape[0], other.shape[1]))
|
||||
|
||||
# For 1D x 2D: treat 1D as row vector
|
||||
if self.ndim == 1 and other.ndim == 2 and isinstance(other, tensor):
|
||||
if self.shape[0] != other.shape[0]:
|
||||
raise ValueError(f"Shape mismatch: ({self.shape[0]},) @ {other.shape}")
|
||||
|
||||
result_values = []
|
||||
for j in range(other.shape[1]):
|
||||
dot_sum = sum(self.get_scalar(k) * other.get_scalar(k, j) for k in range(self.shape[0]))
|
||||
result_values.append(dot_sum)
|
||||
|
||||
return tensor(tuple(result_values), (other.shape[1],))
|
||||
|
||||
# For 2D x 1D: treat 1D as column vector
|
||||
if self.ndim == 2 and other.ndim == 1:
|
||||
if self.shape[1] != other.shape[0]:
|
||||
raise ValueError(f"Shape mismatch: {self.shape} @ ({other.shape[0]},)")
|
||||
|
||||
result_values = []
|
||||
|
||||
if isinstance(other, vector):
|
||||
for i in range(self.shape[0]):
|
||||
dot_sum = value(0)
|
||||
for k in range(self.shape[1]):
|
||||
dot_sum = dot_sum + self.get_scalar(i, k) * other.get_scalar(k)
|
||||
result_values.append(dot_sum)
|
||||
else:
|
||||
for i in range(self.shape[0]):
|
||||
dot_sum = value(0)
|
||||
for k in range(self.shape[1]):
|
||||
dot_sum = dot_sum + self.get_scalar(i, k) * other.get_scalar(k)
|
||||
result_values.append(dot_sum)
|
||||
|
||||
return tensor(tuple(result_values), (self.shape[0],))
|
||||
|
||||
raise NotImplementedError(f"matmul not implemented for shapes {self.ndim}D @ {other.ndim}D")
|
||||
|
||||
def __matmul__(self, other: 'tensor[TNum] | vector[TNum]') -> 'TNum | value[TNum] | tensor[TNum]':
|
||||
"""Matrix multiplication operator (@)."""
|
||||
return self.matmul(other)
|
||||
|
||||
def __rmatmul__(self, other: 'tensor[TNum] | vector[TNum]') -> 'TNum | value[TNum] | tensor[TNum]':
|
||||
"""Right matrix multiplication operator."""
|
||||
if isinstance(other, tensor):
|
||||
return other.matmul(self)
|
||||
return NotImplemented
|
||||
|
||||
def sum(self, axis: int | Sequence[int] | None = None, keepdims: bool = False) -> TNum | value[TNum] | 'tensor[TNum]':
|
||||
"""Sum all or along specified axis/axes.
|
||||
|
||||
Arguments:
|
||||
axis: Axis or tuple of axes along which to sum. If None, sums all elements.
|
||||
keepdims: If True, keep reduced dimensions as size 1.
|
||||
|
||||
Returns:
|
||||
Scalar or tensor with reduced dimension(s).
|
||||
"""
|
||||
if axis is None:
|
||||
result = mixed_sum(self.values)
|
||||
if keepdims:
|
||||
# Return tensor with all dimensions set to 1
|
||||
new_shape = [1 for _ in self.shape]
|
||||
return tensor((result,), new_shape)
|
||||
return result
|
||||
|
||||
# Handle single axis (convert to tuple for uniform processing)
|
||||
if isinstance(axis, int):
|
||||
axes: tuple[int, ...] = (axis,)
|
||||
else:
|
||||
axes = tuple(axis)
|
||||
|
||||
# Validate and normalize axes
|
||||
normalized_axes: list[int] = []
|
||||
for ax in axes:
|
||||
if not (0 <= ax < self.ndim):
|
||||
raise ValueError(f"Axis {ax} is out of bounds for tensor of rank {self.ndim}")
|
||||
if ax not in normalized_axes:
|
||||
normalized_axes.append(ax)
|
||||
|
||||
# Sort axes in descending order for easier dimension removal
|
||||
normalized_axes = sorted(set(normalized_axes), reverse=True)
|
||||
|
||||
# Sum along specified axes
|
||||
new_shape = list(self.shape)
|
||||
for ax in normalized_axes:
|
||||
new_shape.pop(ax)
|
||||
|
||||
if not new_shape:
|
||||
# All axes summed - return scalar
|
||||
return mixed_sum(self.values)
|
||||
|
||||
new_size = 1
|
||||
for dim in new_shape:
|
||||
new_size *= dim
|
||||
|
||||
new_values: list[TNum | value[TNum]] = [self.values[0]] * new_size
|
||||
new_v_mask: list[bool] = [False] * new_size
|
||||
|
||||
for old_idx in range(len(self.values)):
|
||||
old_indices = list(self._get_indices_from_flat(old_idx))
|
||||
|
||||
# Build new indices by removing summed axes
|
||||
new_indices: list[int] = []
|
||||
for i, idx in enumerate(old_indices):
|
||||
if i not in normalized_axes:
|
||||
new_indices.append(idx)
|
||||
|
||||
# Compute flat index in new shape
|
||||
new_flat_idx = 0
|
||||
stride = 1
|
||||
for i in range(len(new_shape) - 1, -1, -1):
|
||||
new_flat_idx += new_indices[i] * stride
|
||||
stride *= new_shape[i]
|
||||
|
||||
if new_v_mask[new_flat_idx]:
|
||||
new_values[new_flat_idx] = new_values[new_flat_idx] + self.values[old_idx]
|
||||
else:
|
||||
new_values[new_flat_idx] = self.values[old_idx]
|
||||
new_v_mask[new_flat_idx] = True
|
||||
|
||||
if keepdims:
|
||||
# Restore reduced dimensions as size 1
|
||||
full_shape = list(self.shape)
|
||||
for ax in normalized_axes:
|
||||
full_shape[ax] = 1
|
||||
return tensor(new_values, tuple(full_shape))
|
||||
|
||||
if not new_shape:
|
||||
return new_values[0]
|
||||
return tensor(new_values, tuple(new_shape))
|
||||
|
||||
def mean(self, axis: int | None = None) -> Any:
|
||||
"""Calculate mean along axis or overall.
|
||||
|
||||
Arguments:
|
||||
axis: Axis along which to compute mean. If None, computes overall mean.
|
||||
|
||||
Returns:
|
||||
Scalar or tensor with reduced dimension.
|
||||
"""
|
||||
if axis is None:
|
||||
total_sum: Any = mixed_sum(self.values)
|
||||
return total_sum / self.size()
|
||||
|
||||
sum_result: Any = self.sum(axis)
|
||||
axis_size = self.shape[axis]
|
||||
|
||||
if isinstance(sum_result, tensor):
|
||||
return sum_result / axis_size
|
||||
else:
|
||||
return sum_result / axis_size
|
||||
|
||||
def map(self, func: Callable[[Any], value[U] | U]) -> 'tensor[U]':
|
||||
"""Apply a function to each element.
|
||||
|
||||
Arguments:
|
||||
func: Function to apply to each element.
|
||||
|
||||
Returns:
|
||||
A new tensor with the function applied element-wise.
|
||||
"""
|
||||
result_vals = tuple(func(v) for v in self.values)
|
||||
return tensor(result_vals, self.shape)
|
||||
|
||||
def homogenize(self) -> 'tensor[TNum]':
|
||||
"""Convert all elements to copapy values if any element is a copapy value."""
|
||||
if any(isinstance(val, value) for val in self.values):
|
||||
homogenized: tuple[value[Any], ...] = tuple(value(val) if not isinstance(val, value) else val for val in self.values)
|
||||
return tensor(homogenized, self.shape)
|
||||
return self
|
||||
|
||||
@property
|
||||
def T(self) -> 'tensor[TNum]':
|
||||
"""Transpose all axes (equivalent to transpose() with no args)."""
|
||||
return self.transpose()
|
||||
|
||||
|
||||
def zeros(shape: Sequence[int] | int) -> tensor[int]:
|
||||
"""Create a zero tensor of given shape."""
|
||||
if isinstance(shape, int):
|
||||
shape = (shape,)
|
||||
|
||||
size = 1
|
||||
for dim in shape:
|
||||
size *= dim
|
||||
|
||||
return tensor([0] * size, tuple(shape))
|
||||
|
||||
|
||||
def ones(shape: Sequence[int] | int) -> tensor[int]:
|
||||
"""Create a tensor of ones with given shape."""
|
||||
if isinstance(shape, int):
|
||||
shape = (shape,)
|
||||
|
||||
size = 1
|
||||
for dim in shape:
|
||||
size *= dim
|
||||
|
||||
return tensor([1] * size, tuple(shape))
|
||||
|
||||
|
||||
def arange(start: int | float, stop: int | float | None = None,
|
||||
step: int | float = 1) -> tensor[int] | tensor[float]:
|
||||
"""Create a tensor with evenly spaced values.
|
||||
|
||||
Arguments:
|
||||
start: Start value (or stop if stop is None).
|
||||
stop: Stop value (exclusive).
|
||||
step: Step between values.
|
||||
|
||||
Returns:
|
||||
A 1D tensor.
|
||||
"""
|
||||
if stop is None:
|
||||
stop = start
|
||||
start = 0
|
||||
|
||||
# Determine type
|
||||
values_list: list[value[Any]] = []
|
||||
current = start
|
||||
if step > 0:
|
||||
while current < stop:
|
||||
values_list.append(value(current))
|
||||
current += step
|
||||
elif step < 0:
|
||||
while current > stop:
|
||||
values_list.append(value(current))
|
||||
current += step
|
||||
else:
|
||||
raise ValueError("step cannot be zero")
|
||||
|
||||
return tensor(tuple(values_list), (len(values_list),))
|
||||
|
||||
|
||||
def eye(rows: int, cols: int | None = None) -> tensor[int]:
|
||||
"""Create an identity tensor with ones on diagonal.
|
||||
|
||||
Arguments:
|
||||
rows: Number of rows.
|
||||
cols: Number of columns (defaults to rows).
|
||||
|
||||
Returns:
|
||||
A 2D identity tensor.
|
||||
"""
|
||||
if cols is None:
|
||||
cols = rows
|
||||
|
||||
values_list: list[value[int]] = []
|
||||
|
||||
for i in range(rows):
|
||||
for j in range(cols):
|
||||
if i == j:
|
||||
values_list.append(value(1))
|
||||
else:
|
||||
values_list.append(value(0))
|
||||
|
||||
return tensor(tuple(values_list), (rows, cols))
|
||||
|
||||
|
||||
def identity(size: int) -> tensor[int]:
|
||||
"""Create a square identity tensor.
|
||||
|
||||
Arguments:
|
||||
size: Size of the square tensor.
|
||||
|
||||
Returns:
|
||||
A square 2D identity tensor.
|
||||
"""
|
||||
return eye(size, size)
|
||||
|
||||
|
||||
@overload
|
||||
def diagonal(vec: 'tensor[int] | vector[int]') -> tensor[int]: ...
|
||||
@overload
|
||||
def diagonal(vec: 'tensor[float] | vector[float]') -> tensor[float]: ...
|
||||
def diagonal(vec: 'tensor[Any] | vector[Any]') -> 'tensor[Any]':
|
||||
"""Create a diagonal tensor from a 1D tensor.
|
||||
|
||||
Arguments:
|
||||
vec: A 1D tensor with values to place on the diagonal.
|
||||
|
||||
Returns:
|
||||
A 2D tensor with the input values on the diagonal and zeros elsewhere.
|
||||
"""
|
||||
if vec.ndim != 1:
|
||||
raise ValueError(f"Input must be 1D, got {vec.ndim}D")
|
||||
|
||||
size = len(vec)
|
||||
values_list: list[Any] = []
|
||||
|
||||
for i in range(size):
|
||||
for j in range(size):
|
||||
if i == j:
|
||||
values_list.append(vec[i])
|
||||
else:
|
||||
values_list.append(value(0))
|
||||
|
||||
return tensor(tuple(values_list), (size, size))
|
||||
|
|
@ -1,8 +1,9 @@
|
|||
from . import value
|
||||
from ._mixed import mixed_sum, mixed_homogenize
|
||||
from typing import Sequence, TypeVar, Iterable, Any, overload, TypeAlias, Callable, Iterator, Generic
|
||||
from typing import Sequence, TypeVar, Iterable, Any, overload, TypeAlias, Callable, Iterator
|
||||
import copapy as cp
|
||||
from ._helper_types import TNum
|
||||
from ._basic_types import ArrayType
|
||||
|
||||
#VecNumLike: TypeAlias = 'vector[int] | vector[float] | value[int] | value[float] | int | float | bool'
|
||||
VecNumLike: TypeAlias = 'vector[Any] | value[Any] | int | float | bool'
|
||||
|
|
@ -13,8 +14,13 @@ U = TypeVar("U", int, float)
|
|||
epsilon = 1e-20
|
||||
|
||||
|
||||
class vector(Generic[TNum]):
|
||||
class vector(ArrayType[TNum]):
|
||||
"""Mathematical vector class supporting basic operations and interactions with values.
|
||||
|
||||
Attributes:
|
||||
values (tuple[value[TNum] | TNum, ...]): The elements of the vector.
|
||||
ndim (int): Number of dimensions (always 1 for vector).
|
||||
shape (tuple[int, ...]): Shape of the vector as a tuple.
|
||||
"""
|
||||
def __init__(self, values: Iterable[TNum | value[TNum]]):
|
||||
"""Create a vector with given values.
|
||||
|
|
@ -23,6 +29,8 @@ class vector(Generic[TNum]):
|
|||
values: iterable of constant values
|
||||
"""
|
||||
self.values: tuple[value[TNum] | TNum, ...] = tuple(values)
|
||||
self.ndim: int = 1
|
||||
self.shape: tuple[int, ...] = (len(self.values),)
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"vector({self.values})"
|
||||
|
|
@ -45,6 +53,10 @@ class vector(Generic[TNum]):
|
|||
def __iter__(self) -> Iterator[value[TNum] | TNum]:
|
||||
return iter(self.values)
|
||||
|
||||
def get_scalar(self, index: int) -> TNum | value[TNum]:
|
||||
"""Get a single scalar value from the vector."""
|
||||
return self.values[index]
|
||||
|
||||
@overload
|
||||
def __add__(self: 'vector[int]', other: VecFloatLike) -> 'vector[float]': ...
|
||||
@overload
|
||||
|
|
@ -213,7 +225,7 @@ class vector(Generic[TNum]):
|
|||
a3 * b1 - a1 * b3,
|
||||
a1 * b2 - a2 * b1
|
||||
])
|
||||
|
||||
|
||||
def __gt__(self, other: VecNumLike) -> 'vector[int]':
|
||||
if isinstance(other, vector):
|
||||
assert len(self.values) == len(other.values)
|
||||
|
|
@ -267,11 +279,7 @@ class vector(Generic[TNum]):
|
|||
return vector(a != other for a in self.values)
|
||||
o = value(other) # Make sure a single constant is allocated
|
||||
return vector(a != o if isinstance(a, value) else a != other for a in self.values)
|
||||
|
||||
@property
|
||||
def shape(self) -> tuple[int]:
|
||||
"""Return the shape of the vector as (length,)."""
|
||||
return (len(self.values),)
|
||||
|
||||
|
||||
@overload
|
||||
def sum(self: 'vector[int]') -> int | value[int]: ...
|
||||
|
|
@ -299,15 +307,15 @@ class vector(Generic[TNum]):
|
|||
|
||||
def map(self, func: Callable[[Any], value[U] | U]) -> 'vector[U]':
|
||||
"""Applies a function to each element of the vector and returns a new vector.
|
||||
|
||||
|
||||
Arguments:
|
||||
func: A function that takes a single argument.
|
||||
|
||||
|
||||
Returns:
|
||||
A new vector with the function applied to each element.
|
||||
"""
|
||||
return vector(func(x) for x in self.values)
|
||||
|
||||
|
||||
def _map2(self, other: VecNumLike, func: Callable[[Any, Any], value[int] | value[float]]) -> 'vector[Any]':
|
||||
if isinstance(other, vector):
|
||||
assert len(self.values) == len(other.values)
|
||||
|
|
@ -320,11 +328,11 @@ class vector(Generic[TNum]):
|
|||
|
||||
def cross_product(v1: vector[float], v2: vector[float]) -> vector[float]:
|
||||
"""Calculate the cross product of two 3D vectors.
|
||||
|
||||
|
||||
Arguments:
|
||||
v1: First 3D vector.
|
||||
v2: Second 3D vector.
|
||||
|
||||
|
||||
Returns:
|
||||
The cross product vector.
|
||||
"""
|
||||
|
|
@ -333,11 +341,11 @@ def cross_product(v1: vector[float], v2: vector[float]) -> vector[float]:
|
|||
|
||||
def dot_product(v1: vector[float], v2: vector[float]) -> 'float | value[float]':
|
||||
"""Calculate the dot product of two vectors.
|
||||
|
||||
|
||||
Arguments:
|
||||
v1: First vector.
|
||||
v2: Second vector.
|
||||
|
||||
|
||||
Returns:
|
||||
The dot product.
|
||||
"""
|
||||
|
|
@ -346,11 +354,11 @@ def dot_product(v1: vector[float], v2: vector[float]) -> 'float | value[float]':
|
|||
|
||||
def distance(v1: vector[float], v2: vector[float]) -> 'float | value[float]':
|
||||
"""Calculate the Euclidean distance between two vectors.
|
||||
|
||||
|
||||
Arguments:
|
||||
v1: First vector.
|
||||
v2: Second vector.
|
||||
|
||||
|
||||
Returns:
|
||||
The Euclidean distance.
|
||||
"""
|
||||
|
|
@ -360,11 +368,11 @@ def distance(v1: vector[float], v2: vector[float]) -> 'float | value[float]':
|
|||
|
||||
def scalar_projection(v1: vector[float], v2: vector[float]) -> 'float | value[float]':
|
||||
"""Calculate the scalar projection of v1 onto v2.
|
||||
|
||||
|
||||
Arguments:
|
||||
v1: First vector.
|
||||
v2: Second vector.
|
||||
|
||||
|
||||
Returns:
|
||||
The scalar projection.
|
||||
"""
|
||||
|
|
@ -375,11 +383,11 @@ def scalar_projection(v1: vector[float], v2: vector[float]) -> 'float | value[fl
|
|||
|
||||
def vector_projection(v1: vector[float], v2: vector[float]) -> vector[float]:
|
||||
"""Calculate the vector projection of v1 onto v2.
|
||||
|
||||
|
||||
Arguments:
|
||||
v1: First vector.
|
||||
v2: Second vector.
|
||||
|
||||
|
||||
Returns:
|
||||
The projected vector.
|
||||
"""
|
||||
|
|
@ -391,11 +399,11 @@ def vector_projection(v1: vector[float], v2: vector[float]) -> vector[float]:
|
|||
|
||||
def angle_between(v1: vector[float], v2: vector[float]) -> 'float | value[float]':
|
||||
"""Calculate the angle in radians between two vectors.
|
||||
|
||||
|
||||
Arguments:
|
||||
v1: First vector.
|
||||
v2: Second vector.
|
||||
|
||||
|
||||
Returns:
|
||||
The angle in radians.
|
||||
"""
|
||||
|
|
@ -408,12 +416,12 @@ def angle_between(v1: vector[float], v2: vector[float]) -> 'float | value[float]
|
|||
|
||||
def rotate_vector(v: vector[float], axis: vector[float], angle: 'float | value[float]') -> vector[float]:
|
||||
"""Rotate vector v around a given axis by a specified angle using Rodrigues' rotation formula.
|
||||
|
||||
|
||||
Arguments:
|
||||
v: The 3D vector to be rotated.
|
||||
axis: A 3D vector defining the axis of rotation.
|
||||
angle: The angle of rotation in radians.
|
||||
|
||||
|
||||
Returns:
|
||||
The rotated vector.
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -83,6 +83,15 @@ def get_cast(type1: str, type2: str, type_out: str) -> str:
|
|||
"""
|
||||
|
||||
|
||||
@norm_indent
|
||||
def get_neg(type1: str) -> str:
|
||||
return f"""
|
||||
STENCIL void neg_{type1}({type1} arg1) {{
|
||||
result_{type1}(-arg1);
|
||||
}}
|
||||
"""
|
||||
|
||||
|
||||
@norm_indent
|
||||
def get_func1(func_name: str, type1: str) -> str:
|
||||
return f"""
|
||||
|
|
@ -249,6 +258,9 @@ if __name__ == "__main__":
|
|||
for fn, t1 in permutate(fnames, types):
|
||||
code += get_func1(fn, t1)
|
||||
|
||||
for t in types:
|
||||
code += get_neg(t)
|
||||
|
||||
fnames = ['sqrt', 'exp', 'log', 'sin', 'cos', 'tan', 'asin', 'acos', 'atan']
|
||||
for fn, t1 in permutate(fnames, types):
|
||||
code += get_math_func1(fn + 'f', t1, fn)
|
||||
|
|
|
|||
|
|
@ -4,26 +4,24 @@ import pytest
|
|||
|
||||
def test_matrix_init():
|
||||
"""Test basic matrix initialization"""
|
||||
m1 = cp.matrix([[1, 2, 3], [4, 5, 6]])
|
||||
assert m1.rows == 2
|
||||
assert m1.cols == 3
|
||||
m1 = cp.tensor([[1, 2, 3], [4, 5, 6]])
|
||||
assert m1.shape == (2, 3)
|
||||
assert m1[0] == (1, 2, 3)
|
||||
assert m1[1] == (4, 5, 6)
|
||||
|
||||
|
||||
def test_matrix_with_variables():
|
||||
"""Test matrix initialization with variables"""
|
||||
m1 = cp.matrix([[cp.value(1), 2], [3, cp.value(4)]])
|
||||
assert m1.rows == 2
|
||||
assert m1.cols == 2
|
||||
assert isinstance(m1[0][0], cp.value)
|
||||
assert isinstance(m1[1][1], cp.value)
|
||||
m1 = cp.tensor([[cp.value(1), 2], [3, cp.value(4)]])
|
||||
assert m1.shape == (2, 2)
|
||||
assert isinstance(m1[0][0], cp.tensor)
|
||||
assert isinstance(m1[1][1], cp.tensor)
|
||||
|
||||
|
||||
def test_matrix_addition():
|
||||
"""Test matrix addition"""
|
||||
m1 = cp.matrix([[1, 2], [3, 4]])
|
||||
m2 = cp.matrix([[5, 6], [7, 8]])
|
||||
m1 = cp.tensor([[1, 2], [3, 4]])
|
||||
m2 = cp.tensor([[5, 6], [7, 8]])
|
||||
m3 = m1 + m2
|
||||
|
||||
assert m3[0] == (6, 8)
|
||||
|
|
@ -32,7 +30,7 @@ def test_matrix_addition():
|
|||
|
||||
def test_matrix_scalar_addition():
|
||||
"""Test matrix addition with scalar"""
|
||||
m1 = cp.matrix([[1, 2], [3, 4]])
|
||||
m1 = cp.tensor([[1, 2], [3, 4]])
|
||||
m2 = m1 + 5
|
||||
|
||||
assert m2[0] == (6, 7)
|
||||
|
|
@ -41,8 +39,8 @@ def test_matrix_scalar_addition():
|
|||
|
||||
def test_matrix_subtraction():
|
||||
"""Test matrix subtraction"""
|
||||
m1 = cp.matrix([[5, 6], [7, 8]])
|
||||
m2 = cp.matrix([[1, 2], [3, 4]])
|
||||
m1 = cp.tensor([[5, 6], [7, 8]])
|
||||
m2 = cp.tensor([[1, 2], [3, 4]])
|
||||
m3 = m1 - m2
|
||||
|
||||
assert m3[0] == (4, 4)
|
||||
|
|
@ -51,7 +49,7 @@ def test_matrix_subtraction():
|
|||
|
||||
def test_matrix_scalar_subtraction():
|
||||
"""Test matrix subtraction with scalar"""
|
||||
m1 = cp.matrix([[5, 6], [7, 8]])
|
||||
m1 = cp.tensor([[5, 6], [7, 8]])
|
||||
m2 = m1 - 2
|
||||
|
||||
assert m2[0] == (3, 4)
|
||||
|
|
@ -60,17 +58,19 @@ def test_matrix_scalar_subtraction():
|
|||
|
||||
def test_matrix_negation():
|
||||
"""Test matrix negation"""
|
||||
m1 = cp.matrix([[1, 2], [3, 4]])
|
||||
m1 = cp.tensor([[1, 2], [3, 4]])
|
||||
m2 = -m1
|
||||
|
||||
assert m1[0] == (1, 2)
|
||||
assert m1[1] == (3, 4)
|
||||
assert m2[0] == (-1, -2)
|
||||
assert m2[1] == (-3, -4)
|
||||
|
||||
|
||||
def test_matrix_element_wise_multiplication():
|
||||
"""Test element-wise matrix multiplication"""
|
||||
m1 = cp.matrix([[1, 2], [3, 4]])
|
||||
m2 = cp.matrix([[5, 6], [7, 8]])
|
||||
m1 = cp.tensor([[1, 2], [3, 4]])
|
||||
m2 = cp.tensor([[5, 6], [7, 8]])
|
||||
m3 = m1 * m2
|
||||
|
||||
assert m3[0] == (5, 12)
|
||||
|
|
@ -79,7 +79,7 @@ def test_matrix_element_wise_multiplication():
|
|||
|
||||
def test_matrix_scalar_multiplication():
|
||||
"""Test matrix multiplication with scalar"""
|
||||
m1 = cp.matrix([[1, 2], [3, 4]])
|
||||
m1 = cp.tensor([[1, 2], [3, 4]])
|
||||
m2 = m1 * 3
|
||||
|
||||
assert m2[0] == (3, 6)
|
||||
|
|
@ -88,8 +88,8 @@ def test_matrix_scalar_multiplication():
|
|||
|
||||
def test_matrix_element_wise_division():
|
||||
"""Test element-wise matrix division"""
|
||||
m1 = cp.matrix([[6.0, 8.0], [12.0, 16.0]])
|
||||
m2 = cp.matrix([[2.0, 2.0], [3.0, 4.0]])
|
||||
m1 = cp.tensor([[6.0, 8.0], [12.0, 16.0]])
|
||||
m2 = cp.tensor([[2.0, 2.0], [3.0, 4.0]])
|
||||
m3 = m1 / m2
|
||||
|
||||
assert m3[0][0] == pytest.approx(3.0) # pyright: ignore[reportUnknownMemberType]
|
||||
|
|
@ -100,7 +100,7 @@ def test_matrix_element_wise_division():
|
|||
|
||||
def test_matrix_scalar_division():
|
||||
"""Test matrix division by scalar"""
|
||||
m1 = cp.matrix([[6.0, 8.0], [12.0, 16.0]])
|
||||
m1 = cp.tensor([[6.0, 8.0], [12.0, 16.0]])
|
||||
m2 = m1 / 2.0
|
||||
|
||||
assert list(m2[0]) == pytest.approx((3.0, 4.0)) # pyright: ignore[reportUnknownMemberType]
|
||||
|
|
@ -109,25 +109,24 @@ def test_matrix_scalar_division():
|
|||
|
||||
def test_matrix_vector_multiplication():
|
||||
"""Test matrix-vector multiplication using @ operator"""
|
||||
m = cp.matrix([[1, 2, 3], [4, 5, 6]])
|
||||
m = cp.tensor([[1, 2, 3], [4, 5, 6]])
|
||||
v = cp.vector([7, 8, 9])
|
||||
result = m @ v
|
||||
|
||||
assert isinstance(result, cp.vector)
|
||||
assert isinstance(result, cp.tensor)
|
||||
assert len(result.values) == 2
|
||||
assert result.values[0] == 1*7 + 2*8 + 3*9
|
||||
assert result.values[1] == 4*7 + 5*8 + 6*9
|
||||
assert result[0] == 1*7 + 2*8 + 3*9
|
||||
assert result[1] == 4*7 + 5*8 + 6*9
|
||||
|
||||
|
||||
def test_matrix_matrix_multiplication():
|
||||
"""Test matrix-matrix multiplication using @ operator"""
|
||||
m1 = cp.matrix([[1, 2], [3, 4]])
|
||||
m2 = cp.matrix([[5, 6], [7, 8]])
|
||||
m1 = cp.tensor([[1, 2], [3, 4]])
|
||||
m2 = cp.tensor([[5, 6], [7, 8]])
|
||||
result = m1 @ m2
|
||||
|
||||
assert isinstance(result, cp.matrix)
|
||||
assert result.rows == 2
|
||||
assert result.cols == 2
|
||||
assert isinstance(result, cp.tensor)
|
||||
assert result.shape == (2, 2)
|
||||
assert result[0][0] == 1*5 + 2*7
|
||||
assert result[0][1] == 1*6 + 2*8
|
||||
assert result[1][0] == 3*5 + 4*7
|
||||
|
|
@ -136,11 +135,10 @@ def test_matrix_matrix_multiplication():
|
|||
|
||||
def test_matrix_transpose():
|
||||
"""Test matrix transpose"""
|
||||
m = cp.matrix([[1, 2, 3], [4, 5, 6]])
|
||||
m = cp.tensor([[1, 2, 3], [4, 5, 6]])
|
||||
mt = m.transpose()
|
||||
|
||||
assert mt.rows == 3
|
||||
assert mt.cols == 2
|
||||
assert mt.shape == (3, 2)
|
||||
assert mt[0] == (1, 4)
|
||||
assert mt[1] == (2, 5)
|
||||
assert mt[2] == (3, 6)
|
||||
|
|
@ -148,35 +146,34 @@ def test_matrix_transpose():
|
|||
|
||||
def test_matrix_transpose_property():
|
||||
"""Test matrix transpose using .T property"""
|
||||
m = cp.matrix([[1, 2, 3], [4, 5, 6]])
|
||||
m = cp.tensor([[1, 2, 3], [4, 5, 6]])
|
||||
mt = m.T
|
||||
|
||||
assert mt.rows == 3
|
||||
assert mt.cols == 2
|
||||
assert mt.shape == (3, 2)
|
||||
assert mt[0] == (1, 4)
|
||||
|
||||
|
||||
def test_matrix_row_access():
|
||||
"""Test getting a row as a vector"""
|
||||
m = cp.matrix([[1, 2, 3], [4, 5, 6]])
|
||||
row0 = m.row(0)
|
||||
m = cp.tensor([[1, 2, 3], [4, 5, 6]])
|
||||
row0 = m[0]
|
||||
|
||||
assert isinstance(row0, cp.vector)
|
||||
assert isinstance(row0, cp.tensor)
|
||||
assert row0.values == (1, 2, 3)
|
||||
|
||||
|
||||
def test_matrix_col_access():
|
||||
"""Test getting a column as a vector"""
|
||||
m = cp.matrix([[1, 2, 3], [4, 5, 6]])
|
||||
col1 = m.col(1)
|
||||
m = cp.tensor([[1, 2, 3], [4, 5, 6]])
|
||||
col1 = m[:, 1]
|
||||
|
||||
assert isinstance(col1, cp.vector)
|
||||
assert col1.values == (2, 5)
|
||||
assert isinstance(col1, cp.tensor)
|
||||
assert col1 == (2, 5)
|
||||
|
||||
|
||||
def test_matrix_trace():
|
||||
"""Test matrix trace (sum of diagonal elements)"""
|
||||
m = cp.matrix([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
|
||||
m = cp.tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
|
||||
trace = m.trace()
|
||||
|
||||
assert trace == 1 + 5 + 9
|
||||
|
|
@ -184,7 +181,7 @@ def test_matrix_trace():
|
|||
|
||||
def test_matrix_sum():
|
||||
"""Test sum of all matrix elements"""
|
||||
m = cp.matrix([[1, 2, 3], [4, 5, 6]])
|
||||
m = cp.tensor([[1, 2, 3], [4, 5, 6]])
|
||||
total = m.sum()
|
||||
|
||||
assert total == 1 + 2 + 3 + 4 + 5 + 6
|
||||
|
|
@ -192,7 +189,7 @@ def test_matrix_sum():
|
|||
|
||||
def test_matrix_map():
|
||||
"""Test mapping a function over matrix elements"""
|
||||
m = cp.matrix([[1, 2], [3, 4]])
|
||||
m = cp.tensor([[1, 2], [3, 4]])
|
||||
m_doubled = m.map(lambda x: x * 2)
|
||||
|
||||
assert m_doubled[0] == (2, 4)
|
||||
|
|
@ -201,20 +198,19 @@ def test_matrix_map():
|
|||
|
||||
def test_matrix_homogenize():
|
||||
"""Test homogenizing matrix (converting to all variables)"""
|
||||
m = cp.matrix([[1, cp.value(2)], [3, 4]])
|
||||
m = cp.tensor([[1, cp.value(2)], [3, 4]])
|
||||
m_homo = m.homogenize()
|
||||
|
||||
for row in m_homo:
|
||||
for elem in row:
|
||||
assert isinstance(elem, cp.value)
|
||||
assert isinstance(elem, cp.tensor) and elem.ndim == 0
|
||||
|
||||
|
||||
def test_identity_matrix():
|
||||
"""Test identity matrix creation"""
|
||||
m = cp.identity(3)
|
||||
|
||||
assert m.rows == 3
|
||||
assert m.cols == 3
|
||||
assert m.shape == (3, 3)
|
||||
assert m[0] == (1, 0, 0)
|
||||
assert m[1] == (0, 1, 0)
|
||||
assert m[2] == (0, 0, 1)
|
||||
|
|
@ -222,20 +218,18 @@ def test_identity_matrix():
|
|||
|
||||
def test_zeros_matrix():
|
||||
"""Test zeros matrix creation"""
|
||||
m = cp.zeros(2, 3)
|
||||
m = cp.zeros([2, 3])
|
||||
|
||||
assert m.rows == 2
|
||||
assert m.cols == 3
|
||||
assert m.shape == (2, 3)
|
||||
assert m[0] == (0, 0, 0)
|
||||
assert m[1] == (0, 0, 0)
|
||||
|
||||
|
||||
def test_ones_matrix():
|
||||
"""Test ones matrix creation"""
|
||||
m = cp.ones(2, 3)
|
||||
m = cp.ones([2, 3])
|
||||
|
||||
assert m.rows == 2
|
||||
assert m.cols == 3
|
||||
assert m.shape == (2, 3)
|
||||
assert m[0] == (1, 1, 1)
|
||||
assert m[1] == (1, 1, 1)
|
||||
|
||||
|
|
@ -245,8 +239,7 @@ def test_diagonal_matrix():
|
|||
v = cp.vector([1, 2, 3])
|
||||
m = cp.diagonal(v)
|
||||
|
||||
assert m.rows == 3
|
||||
assert m.cols == 3
|
||||
assert m.shape == (3, 3)
|
||||
assert m[0] == (1, 0, 0)
|
||||
assert m[1] == (0, 2, 0)
|
||||
assert m[2] == (0, 0, 3)
|
||||
|
|
@ -254,7 +247,7 @@ def test_diagonal_matrix():
|
|||
|
||||
def test_matrix_with_variables_compiled():
|
||||
"""Test matrix operations with variables in compilation"""
|
||||
m = cp.matrix([[cp.value(1), 2], [3, cp.value(4)]])
|
||||
m = cp.tensor([[cp.value(1), 2], [3, cp.value(4)]])
|
||||
v = cp.vector([cp.value(5), 6])
|
||||
result = m @ v
|
||||
|
||||
|
|
|
|||
|
|
@ -0,0 +1,257 @@
|
|||
#!/usr/bin/env python3
|
||||
"""Basic tests for the tensor class."""
|
||||
|
||||
import copapy as cp
|
||||
|
||||
def test_tensor_basic():
|
||||
# Test 1: Create a scalar tensor
|
||||
print("Test 1: Scalar tensor")
|
||||
t0 = cp.tensor(42)
|
||||
print(f"Scalar tensor: {t0}")
|
||||
print(f"Shape: {t0.shape}, ndim: {t0.ndim}")
|
||||
assert t0.shape == ()
|
||||
assert t0 == 42
|
||||
print()
|
||||
|
||||
# Test 2: Create a 1D tensor from list
|
||||
print("Test 2: 1D tensor")
|
||||
t1 = cp.tensor([1, 2, 3, 4, 5])
|
||||
print(f"1D tensor: shape={t1.shape}, ndim={t1.ndim}")
|
||||
print(f"Elements: {[t1[i] for i in range(len(t1))]}")
|
||||
assert t1.shape == (5,)
|
||||
assert t1.ndim == 1
|
||||
assert t1[0] == 1
|
||||
print()
|
||||
|
||||
# Test 3: Create a 2D tensor (matrix)
|
||||
print("Test 3: 2D tensor")
|
||||
t2 = cp.tensor([[1, 2, 3], [4, 5, 6]])
|
||||
print(f"2D tensor: shape={t2.shape}, ndim={t2.ndim}")
|
||||
print(f"Element [0,1]: {t2[0, 1]}")
|
||||
print(f"Row 1: {t2[1]}")
|
||||
assert t2.shape == (2, 3)
|
||||
assert t2.ndim == 2
|
||||
assert t2[0, 1] == 2
|
||||
|
||||
assert t2[1][2] == 6
|
||||
print()
|
||||
|
||||
# Test 4: Create a 3D tensor
|
||||
print("Test 4: 3D tensor")
|
||||
t3 = cp.tensor([[[1, 2], [3, 4]], [[5, 6], [7, 8]]])
|
||||
print(f"3D tensor: shape={t3.shape}, ndim={t3.ndim}")
|
||||
print(f"Element [0,1,0]: {t3[0, 1, 0]}")
|
||||
assert t3.shape == (2, 2, 2)
|
||||
assert t3.ndim == 3
|
||||
assert t3[0, 1, 0] == 3
|
||||
print()
|
||||
|
||||
# Test 6: Broadcasting with scalar
|
||||
print("Test 6: Broadcasting with scalar")
|
||||
t = cp.tensor([1.0, 2.0, 3.0])
|
||||
result = t * 2.0
|
||||
print(f"tensor * 2.0: shape={result.shape}")
|
||||
print(f"Elements: {[result[i] for i in range(len(result))]}")
|
||||
assert result.shape == (3,)
|
||||
assert result[0] == 2.0
|
||||
assert result[1] == 4.0
|
||||
print()
|
||||
|
||||
# Test 6b: Broadcasting with different dimensions
|
||||
print("Test 6b: Broadcasting with different dimensions")
|
||||
# 2D tensor + 1D tensor
|
||||
t2d = cp.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
|
||||
t1d = cp.tensor([10.0, 20.0, 30.0])
|
||||
result_2d_1d = t2d + t1d
|
||||
print(f"2D tensor {t2d.shape} + 1D tensor {t1d.shape} = shape {result_2d_1d.shape}")
|
||||
print(f"Elements: {[[result_2d_1d[i, j] for j in range(3)] for i in range(2)]}")
|
||||
assert result_2d_1d.shape == (2, 3)
|
||||
assert result_2d_1d[0, 0] == 11.0
|
||||
assert result_2d_1d[1, 2] == 36.0
|
||||
|
||||
# 3D tensor + 2D tensor
|
||||
t3d = cp.tensor([[[1.0, 2.0], [3.0, 4.0]], [[5.0, 6.0], [7.0, 8.0]]])
|
||||
t2d_broadcast = cp.tensor([[100.0, 200.0], [300.0, 400.0]])
|
||||
result_3d_2d = t3d + t2d_broadcast
|
||||
print(f"3D tensor {t3d.shape} + 2D tensor {t2d_broadcast.shape} = shape {result_3d_2d.shape}")
|
||||
assert result_3d_2d.shape == (2, 2, 2)
|
||||
assert result_3d_2d[0, 0, 0] == 101.0
|
||||
assert result_3d_2d[1, 1, 1] == 408.0
|
||||
|
||||
# 3D tensor + 1D tensor
|
||||
t1d_broadcast = cp.tensor([1.0, 2.0])
|
||||
result_3d_1d = t3d + t1d_broadcast
|
||||
print(f"3D tensor {t3d.shape} + 1D tensor {t1d_broadcast.shape} = shape {result_3d_1d.shape}")
|
||||
assert result_3d_1d.shape == (2, 2, 2)
|
||||
assert result_3d_1d[0, 0, 0] == 2.0
|
||||
assert result_3d_1d[1, 1, 1] == 10.0
|
||||
print()
|
||||
|
||||
# 3D tensor + vector
|
||||
t1d_broadcast = cp.vector([1.0, 2.0])
|
||||
result_3d_1d = t3d + t1d_broadcast
|
||||
print(f"3D tensor {t3d.shape} + 1D tensor {t1d_broadcast.shape} = shape {result_3d_1d.shape}")
|
||||
assert result_3d_1d.shape == (2, 2, 2)
|
||||
assert result_3d_1d[0, 0, 0] == 2.0
|
||||
assert result_3d_1d[1, 1, 1] == 10.0
|
||||
print()
|
||||
|
||||
# Test 6c: Element-wise operations with different dimensions
|
||||
print("Test 6c: Element-wise operations with different dimensions")
|
||||
a2d = cp.tensor([[1.0, 2.0], [3.0, 4.0]])
|
||||
b2d = cp.tensor([[2.0, 3.0], [4.0, 5.0]])
|
||||
c2d = a2d * b2d
|
||||
print(f"2D * 2D: shape={c2d.shape}")
|
||||
print(f"Elements: {[[c2d[i, j] for j in range(2)] for i in range(2)]}")
|
||||
assert c2d.shape == (2, 2)
|
||||
assert c2d[0, 0] == 2.0
|
||||
assert c2d[1, 1] == 20.0
|
||||
|
||||
# 3D - 2D
|
||||
t3d_sub = cp.tensor([[[10.0, 20.0], [30.0, 40.0]], [[50.0, 60.0], [70.0, 80.0]]])
|
||||
t2d_sub = cp.tensor([[1.0, 2.0], [3.0, 4.0]])
|
||||
result_sub = t3d_sub - t2d_sub
|
||||
print(f"3D - 2D: shape={result_sub.shape}")
|
||||
assert result_sub.shape == (2, 2, 2)
|
||||
assert result_sub[0, 0, 0] == 9.0
|
||||
assert result_sub[1, 1, 1] == 76.0
|
||||
print()
|
||||
|
||||
# Test 7: Reshape
|
||||
print("Test 7: Reshape")
|
||||
t = cp.tensor([1, 2, 3, 4, 5, 6])
|
||||
print(f"Original: shape={t.shape}")
|
||||
t_reshaped = t.reshape(2, 3)
|
||||
print(f"Reshaped to (2, 3): shape={t_reshaped.shape}")
|
||||
print(f"Element [1,2]: {t_reshaped[1, 2]}")
|
||||
assert t_reshaped.shape == (2, 3)
|
||||
assert t_reshaped[1, 2] == 6
|
||||
assert t_reshaped[0, 0] == 1
|
||||
print()
|
||||
|
||||
# Test 8: Flatten
|
||||
print("Test 8: Flatten")
|
||||
t = cp.tensor([[1, 2, 3], [4, 5, 6]])
|
||||
flat = t.flatten()
|
||||
print(f"Original: shape={t.shape}")
|
||||
print(f"Flattened: shape={flat.shape}")
|
||||
print(f"Elements: {[flat[i] for i in range(len(flat))]}")
|
||||
assert flat.shape == (6,)
|
||||
assert flat[0] == 1
|
||||
assert flat[5] == 6
|
||||
print()
|
||||
|
||||
# Test 9: Transpose
|
||||
print("Test 9: Transpose")
|
||||
t = cp.tensor([[1, 2, 3], [4, 5, 6]])
|
||||
print(f"Original: shape={t.shape}")
|
||||
t_t = t.transpose()
|
||||
print(f"Transposed: shape={t_t.shape}")
|
||||
print(f"Element [2,1]: {t_t[2, 1]}")
|
||||
assert t_t.shape == (3, 2)
|
||||
assert t_t[2, 1] == 6
|
||||
print()
|
||||
|
||||
# Test 10: Sum operations
|
||||
print("Test 10: Sum operations")
|
||||
t = cp.tensor([[1, 2, 3], [4, 5, 6]])
|
||||
print(f"Original: shape={t.shape}")
|
||||
total = t.sum()
|
||||
print(f"Sum all: {total}")
|
||||
sum_axis0 = t.sum(axis=0)
|
||||
print(f"Sum along axis 0: shape={sum_axis0.shape}")
|
||||
sum_axis1 = t.sum(axis=1)
|
||||
print(f"Sum along axis 1: shape={sum_axis1.shape}")
|
||||
assert total == 21
|
||||
assert sum_axis0.shape == (3,)
|
||||
assert sum_axis0[0] == 5
|
||||
assert sum_axis1.shape == (2,)
|
||||
assert sum_axis1[1] == 15
|
||||
print()
|
||||
|
||||
# Test 10b: Sum with multiple axes and keepdims
|
||||
print("Test 10b: Sum with multiple axes and keepdims")
|
||||
t3d = cp.tensor([[[1, 2], [3, 4]], [[5, 6], [7, 8]]])
|
||||
print(f"Original 3D tensor: shape={t3d.shape}")
|
||||
|
||||
# Sum along multiple axes
|
||||
sum_axes_0_2 = t3d.sum(axis=(0, 2))
|
||||
print(f"Sum along axes (0, 2): shape={sum_axes_0_2.shape}")
|
||||
assert sum_axes_0_2.shape == (2,), f"Expected (2,), got {sum_axes_0_2.shape}"
|
||||
assert sum_axes_0_2[0] == 1 + 2 + 5 + 6 # Elements from [0,:,*] and [1,:,*]
|
||||
assert sum_axes_0_2[1] == 3 + 4 + 7 + 8
|
||||
print(f"Values: {[sum_axes_0_2[i] for i in range(len(sum_axes_0_2))]}")
|
||||
|
||||
# Sum with keepdims
|
||||
sum_keepdims = t3d.sum(axis=1, keepdims=True)
|
||||
print(f"Sum along axis 1 with keepdims: shape={sum_keepdims.shape}")
|
||||
assert sum_keepdims.shape == (2, 1, 2), f"Expected (2, 1, 2), got {sum_keepdims.shape}"
|
||||
|
||||
# Sum multiple axes with keepdims
|
||||
sum_multi_keepdims = t3d.sum(axis=(0, 2), keepdims=True)
|
||||
print(f"Sum along axes (0, 2) with keepdims: shape={sum_multi_keepdims.shape}")
|
||||
assert sum_multi_keepdims.shape == (1, 2, 1), f"Expected (1, 2, 1), got {sum_multi_keepdims.shape}"
|
||||
|
||||
# Sum all axes with keepdims
|
||||
sum_all_keepdims = t3d.sum(keepdims=True)
|
||||
print(f"Sum all with keepdims: shape={sum_all_keepdims.shape}")
|
||||
assert sum_all_keepdims.shape == (1, 1, 1), f"Expected (1, 1, 1), got {sum_all_keepdims.shape}"
|
||||
assert sum_all_keepdims[0, 0, 0] == 36 # Sum of all elements
|
||||
print()
|
||||
|
||||
# Test 11: Factory functions
|
||||
print("Test 11: Factory functions")
|
||||
z = cp.zeros((2, 3))
|
||||
print(f"zeros((2, 3)): shape={z.shape}")
|
||||
o = cp.ones((3, 2))
|
||||
print(f"ones((3, 2)): shape={o.shape}")
|
||||
e = cp.eye(3)
|
||||
print(f"eye(3): shape={e.shape}")
|
||||
ar = cp.arange(0, 10, 2)
|
||||
print(f"arange(0, 10, 2): shape={ar.shape}")
|
||||
assert z.shape == (2, 3)
|
||||
assert z[0, 0] == 0
|
||||
assert o.shape == (3, 2)
|
||||
assert o[1, 1] == 1
|
||||
print()
|
||||
|
||||
# Test 12: Size and properties
|
||||
print("Test 12: Size and properties")
|
||||
t = cp.tensor([[[1, 2], [3, 4]], [[5, 6], [7, 8]]])
|
||||
print(f"Shape: {t.shape}")
|
||||
print(f"ndim: {t.ndim}")
|
||||
print(f"size: {t.size()}")
|
||||
assert t.shape == (2, 2, 2)
|
||||
assert t.ndim == 3
|
||||
print()
|
||||
|
||||
def test_tensor_slicing():
|
||||
print("Test Numpy-style slicing")
|
||||
t = cp.tensor([[10, 20, 30], [40, 50, 60], [70, 80, 90]])
|
||||
print(f"Original tensor: shape={t.shape}")
|
||||
|
||||
slice1 = t[1]
|
||||
print(f"t[1]: {slice1}, shape={slice1.shape}")
|
||||
assert slice1.shape == (3,)
|
||||
assert slice1[0] == 40
|
||||
|
||||
slice2 = t[:, 2]
|
||||
print(f"t[:, 2]: {slice2}, shape={slice2.shape}")
|
||||
assert slice2.shape == (3,)
|
||||
assert slice2[1] == 60
|
||||
|
||||
slice3 = t[0:2, 1:3]
|
||||
print(f"t[0:2, 1:3]: {slice3}, shape={slice3.shape}")
|
||||
assert slice3.shape == (2, 2)
|
||||
assert slice3[0, 0] == 20
|
||||
|
||||
slice4 = t[-1, :]
|
||||
print(f"t[-1, :]: {slice4}, shape={slice4.shape}")
|
||||
assert slice4.shape == (3,)
|
||||
assert slice4[2] == 90
|
||||
print()
|
||||
|
||||
if __name__ == "__main__":
|
||||
test_tensor_basic()
|
||||
print("All tests completed!")
|
||||
|
||||
|
|
@ -11,20 +11,30 @@ cd musl
|
|||
|
||||
#./configure CFLAGS="-O2 -fno-stack-protector -ffast-math"
|
||||
|
||||
# x86_64
|
||||
sh ../packobjs.sh gcc ld /object_files/musl_objects_x86_64.o
|
||||
|
||||
# x86
|
||||
sh ../packobjs.sh i686-linux-gnu-gcc-13 i686-linux-gnu-ld /object_files/musl_objects_x86.o -fno-pic
|
||||
|
||||
# Arm64
|
||||
sh ../packobjs.sh aarch64-linux-gnu-gcc-13 aarch64-linux-gnu-ld /object_files/musl_objects_arm64.o
|
||||
|
||||
# Armv6
|
||||
sh ../packobjs.sh arm-none-eabi-gcc arm-none-eabi-ld /object_files/musl_objects_armv6.o "-march=armv6 -mfpu=vfp -mfloat-abi=hard -marm"
|
||||
|
||||
# Armv7
|
||||
sh ../packobjs.sh arm-none-eabi-gcc arm-none-eabi-ld /object_files/musl_objects_armv7.o "-march=armv7-a -mfpu=neon-vfpv3 -mfloat-abi=hard -marm"
|
||||
|
||||
# Armv7 Thumb for Cortex-M3..7
|
||||
sh ../packobjs.sh arm-none-eabi-gcc arm-none-eabi-ld /object_files/musl_objects_armv7thumb.o "-march=armv7e-m -mfpu=fpv4-sp-d16 -mfloat-abi=hard -mthumb"
|
||||
|
||||
#sh ../packobjs.sh mips mips-linux-gnu-gcc-13 mips-linux-gnu-ld
|
||||
|
||||
#sh ../packobjs.sh riscv64 riscv64-linux-gnu-gcc-13 riscv64-linux-gnu-ld
|
||||
|
||||
cp ./COPYRIGHT /object_files/
|
||||
|
||||
echo "- clean up..."
|
||||
rm -r ./*
|
||||
cd ..
|
||||
|
|
|
|||
|
|
@ -41,6 +41,11 @@ arm-none-eabi-gcc -march=armv7-a -mfpu=neon-vfpv3 -mfloat-abi=hard -marm $FLAGS
|
|||
LIBGCC=$(arm-none-eabi-gcc -print-libgcc-file-name)
|
||||
arm-none-eabi-ld -r $STMP /object_files/musl_objects_armv7.o $LIBGCC -o $DEST/stencils_armv7_$OPT.o
|
||||
|
||||
# Armv7 Thumb for Cortex-M3..7 hardware fp
|
||||
arm-none-eabi-gcc -march=armv7e-m -mfpu=fpv4-sp-d16 -mfloat-abi=hard -mthumb $FLAGS -$OPT -c $SRC -o $STMP
|
||||
LIBGCC=$(arm-none-eabi-gcc -print-libgcc-file-name)
|
||||
arm-none-eabi-ld -r $STMP /object_files/musl_objects_armv7thumb.o $LIBGCC -o $DEST/stencils_armv7thumb_$OPT.o
|
||||
|
||||
# PowerPC64LE
|
||||
# powerpc64le-linux-gnu-gcc-13 $FLAGS -$OPT -c $SRC -o $DEST/stencils_ppc64le_$OPT.o
|
||||
|
||||
|
|
|
|||
Loading…
Reference in New Issue