mirror of https://github.com/Nonannet/copapy.git
Compare commits
No commits in common. "main" and "v0.0.3-beta" have entirely different histories.
main
...
v0.0.3-bet
|
|
@ -22,11 +22,6 @@ jobs:
|
|||
name: stencil-object-files
|
||||
path: src/copapy/obj/*.o
|
||||
|
||||
- uses: actions/upload-artifact@v4
|
||||
with:
|
||||
name: musl-object-files
|
||||
path: /object_files/*
|
||||
|
||||
build_wheels:
|
||||
if: contains(github.ref, '-beta') == false
|
||||
needs: [build_stencils]
|
||||
|
|
@ -38,25 +33,12 @@ jobs:
|
|||
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
with:
|
||||
fetch-depth: 0
|
||||
fetch-tags: true
|
||||
|
||||
- uses: actions/download-artifact@v4
|
||||
with:
|
||||
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/*
|
||||
path: /object_files/musl_objects_*.*o
|
||||
|
||||
build-package-test:
|
||||
needs: [build_stencils]
|
||||
|
|
@ -47,16 +47,6 @@ 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:
|
||||
|
|
@ -131,7 +121,7 @@ jobs:
|
|||
- uses: actions/upload-artifact@v4
|
||||
if: strategy.job-index == 0
|
||||
with:
|
||||
name: runner-linux-x86_64
|
||||
name: runner-linux
|
||||
path: build/runner/*
|
||||
|
||||
build-arm64:
|
||||
|
|
@ -278,8 +268,10 @@ jobs:
|
|||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
with:
|
||||
fetch-depth: 0
|
||||
fetch-tags: true
|
||||
fetch-depth: 1
|
||||
sparse-checkout: |
|
||||
pyproject.toml
|
||||
tools/get_tag.sh
|
||||
|
||||
- name: Download artifacts
|
||||
uses: actions/download-artifact@v4
|
||||
|
|
@ -297,9 +289,8 @@ jobs:
|
|||
set -v
|
||||
mkdir -p release
|
||||
cp tmp/stencil-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/musl-object-files/* release/
|
||||
cp tmp/runner-linux/coparun release/
|
||||
cp tmp/runner-linux-arm64/coparun release/coparun-aarch64
|
||||
cp tmp/runner-linux-armv6/coparun release/coparun-armv6
|
||||
cp tmp/runner-linux-armv7/coparun release/coparun-armv7
|
||||
|
|
@ -323,8 +314,6 @@ jobs:
|
|||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
with:
|
||||
fetch-tags: true
|
||||
|
||||
- uses: actions/download-artifact@v4
|
||||
with:
|
||||
|
|
@ -350,11 +339,6 @@ jobs:
|
|||
make html
|
||||
touch build/html/.nojekyll
|
||||
|
||||
- name: Upload artifact
|
||||
uses: actions/upload-pages-artifact@v3
|
||||
with:
|
||||
path: docs/build/html
|
||||
|
||||
deploy-docs:
|
||||
if: github.ref == 'refs/heads/main' && github.event_name == 'push'
|
||||
needs: build-docs
|
||||
|
|
|
|||
|
|
@ -29,4 +29,3 @@ docs/source/api
|
|||
core
|
||||
*.log
|
||||
docs/source/start.md
|
||||
/src/copapy/_version.py
|
||||
|
|
|
|||
33
README.md
33
README.md
|
|
@ -1,6 +1,6 @@
|
|||
# Copapy
|
||||
|
||||
Copapy is a Python framework for deterministic, low-latency realtime computation with automatic differentiation support, targeting hardware applications - for example in the fields of robotics, aerospace, SDR, embedded systems and control systems in general.
|
||||
Copapy is a Python framework for deterministic, low-latency realtime computation with automatic differentiation support, targeting hardware applications - for example in the fields of robotics, aerospace, embedded systems and control systems in general.
|
||||
|
||||
GPU frameworks like PyTorch, JAX and TensorFlow jump-started the development in the field of AI. With the right balance of flexibility and performance, they allow for fast iteration of new ideas while still being performant enough to test or even use them in production.
|
||||
|
||||
|
|
@ -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, ARMv6, ARMv7 and AArch64
|
||||
- Optimized machine code for x86_64, AArch64 and ARMv7
|
||||
- Highly portable to new architectures
|
||||
- Small Python package with minimal dependencies and no cross-compile toolchain required
|
||||
|
||||
|
|
@ -38,29 +38,10 @@ 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. It is to note that in this benchmark the copapy case does not move any data between python and the compiled code.
|
||||
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.
|
||||
|
||||
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):
|
||||
|
|
@ -187,11 +168,11 @@ For more complex operations - where inlining is less useful - stencils call a no
|
|||
e: R_X86_64_PLT32 result_float-0x4
|
||||
```
|
||||
|
||||
Unlike stencils, non-stencil functions like `sinf` are not stripped and do not need to be tail-call-optimizable. These functions can be provided as C code and compiled together with the stencils or can be object files like in the case of `sinf` compiled from C and assembly code and merged into the stencil object files. Math functions like `sinf` are currently provided by the MUSL C library, with architecture-specific optimizations.
|
||||
Unlike stencils, non-stencil functions are not stripped and do not need to be tail-call-optimizable.
|
||||
|
||||
Non-stencil functions and constants are stored together with the stencils in an ELF object file for each supported CPU architecture. The required non-stencil functions and constants are bundled during compilation. The compiler includes only the data and code required for a specific Copapy program.
|
||||
Non-stencil functions and constants are stored together with the stencils in an ELF object file for each supported CPU architecture. The required non-stencil functions and constants are bundled during compilation. The compiler includes only the data and code required for the specific program.
|
||||
|
||||
The Copapy compilation process is independent of the actual instruction set. It relies purely on relocation entries and symbol metadata from the ELF file generated by the C compiler.
|
||||
The whole compilation process is independent of the actual instruction set. It relies purely on relocation entries and symbol metadata from the ELF file generated by the C compiler.
|
||||
|
||||
## Developer Guide
|
||||
|
||||
|
|
@ -253,4 +234,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, 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.
|
||||
[^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.
|
||||
|
|
|
|||
|
|
@ -1,4 +1,4 @@
|
|||
# How it works
|
||||
# Compiler
|
||||
```{toctree}
|
||||
:maxdepth: 1
|
||||
:hidden:
|
||||
|
|
|
|||
|
|
@ -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 example script')
|
||||
parser.add_argument('--input', default='tools/make_example.py', help='Path to input C file')
|
||||
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')
|
||||
|
||||
|
|
|
|||
|
|
@ -12,7 +12,7 @@ def extract_sections(md_text: str) -> dict[str, str]:
|
|||
# regex captures: heading marks (###...), heading text, and the following content
|
||||
pattern = re.compile(
|
||||
r'^(#{1,6})\s+(.*?)\s*$' # heading level + heading text
|
||||
r'(.*?(?:```.*?```.*?)*?)' # section content (lazy)
|
||||
r'(.*?)' # section content (lazy)
|
||||
r'(?=^#{1,6}\s+|\Z)', # stop at next heading or end of file
|
||||
re.MULTILINE | re.DOTALL
|
||||
)
|
||||
|
|
@ -37,9 +37,7 @@ if __name__ == '__main__':
|
|||
readme = extract_sections(f.read())
|
||||
|
||||
with open(os.path.join(build_dir, 'start.md'), 'wt') as f:
|
||||
f.write('\n'.join(f"{s}\n" + readme[s.strip(' #')] for s in [
|
||||
'# Copapy', '## Current state', '## Install', '## Examples',
|
||||
'### Basic example', '### Inverse kinematics', '## License']))
|
||||
f.write('\n'.join(f"# {s}\n" + readme[s] for s in ['Copapy', 'Current state', 'Install', 'License']))
|
||||
|
||||
with open(os.path.join(build_dir, 'compiler.md'), 'wt') as f:
|
||||
f.write('\n'.join(readme[s] for s in ['How it works']))
|
||||
|
|
|
|||
|
|
@ -97,7 +97,7 @@ if __name__ == "__main__":
|
|||
|
||||
write_functions(f, ['*'], 'copapy', title='Vector functions', path_patterns=['*_vectors*'], api_dir=api_dir)
|
||||
|
||||
write_functions(f, ['*'], 'copapy', title='Tensor/Matrix functions', path_patterns=['*_tensors*'], api_dir=api_dir)
|
||||
write_functions(f, ['*'], 'copapy', title='Matrix functions', path_patterns=['*_matrices*'], api_dir=api_dir)
|
||||
|
||||
#write_manual(f, ['NumLike'], title='Types')
|
||||
|
||||
|
|
|
|||
|
|
@ -134,7 +134,6 @@ if __name__ == "__main__":
|
|||
md_code: str = ''
|
||||
|
||||
for function_name, code in functions.items():
|
||||
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:
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
[project]
|
||||
name = "copapy"
|
||||
dynamic = ["version"]
|
||||
version = "0.0.3"
|
||||
authors = [
|
||||
{ name="Nicolas Kruse", email="nicolas.kruse@nonan.net" },
|
||||
]
|
||||
|
|
@ -20,7 +20,7 @@ Homepage = "https://github.com/nonannet/copapy"
|
|||
Issues = "https://github.com/nonannet/copapy/issues"
|
||||
|
||||
[build-system]
|
||||
requires = ["setuptools>=61.0", "setuptools-scm>=8", "wheel"]
|
||||
requires = ["setuptools>=61.0", "wheel"]
|
||||
build-backend = "setuptools.build_meta"
|
||||
|
||||
[tool.setuptools.packages.find]
|
||||
|
|
@ -29,17 +29,6 @@ where = ["src"]
|
|||
[tool.setuptools.package-data]
|
||||
copapy = ["obj/*.o", "py.typed"]
|
||||
|
||||
[tool.setuptools_scm]
|
||||
version_scheme = "post-release"
|
||||
local_scheme = "no-local-version"
|
||||
tag_regex = "^v(?P<version>\\d+\\.\\d+\\.\\d+(?:-beta)?)$"
|
||||
fallback_version = "0.0.0"
|
||||
write_to = "src/copapy/_version.py"
|
||||
write_to_template = '''
|
||||
# generated by setuptools_scm - do not edit
|
||||
__version__ = "{version}"
|
||||
'''
|
||||
|
||||
[project.optional-dependencies]
|
||||
dev = [
|
||||
"ruff",
|
||||
|
|
|
|||
|
|
@ -36,28 +36,22 @@ 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 ._tensors import tensor, zeros, ones, arange, eye, identity, diagonal
|
||||
from ._matrices import matrix, identity, zeros, ones, diagonal, eye
|
||||
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
|
||||
from ._version import __version__ # Run "pip install -e ." to generate _version.py
|
||||
|
||||
|
||||
__all__ = [
|
||||
"__version__",
|
||||
"Target",
|
||||
"NumLike",
|
||||
"value",
|
||||
"generic_sdb",
|
||||
"iif",
|
||||
"vector",
|
||||
"tensor",
|
||||
"matrix",
|
||||
"identity",
|
||||
"zeros",
|
||||
"ones",
|
||||
"diagonal",
|
||||
"arange",
|
||||
"sqrt",
|
||||
"abs",
|
||||
"sin",
|
||||
|
|
|
|||
|
|
@ -1,4 +1,4 @@
|
|||
from . import value, vector, tensor
|
||||
from . import value, vector, matrix
|
||||
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: tensor[Any]) -> tensor[float]: ...
|
||||
@overload
|
||||
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:
|
||||
@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:
|
||||
"""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] | tensor[Any
|
|||
|
||||
if isinstance(y, value):
|
||||
y_set = {y}
|
||||
if isinstance(y, tensor):
|
||||
y_set = {v.get_scalar(0) for v in y.flatten()}
|
||||
if isinstance(y, matrix):
|
||||
y_set = {v for row in y for v in row}
|
||||
else:
|
||||
assert isinstance(y, Sequence) or isinstance(y, vector)
|
||||
y_set = set(y)
|
||||
y_set = {v for v in 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)
|
||||
|
|
@ -89,11 +89,8 @@ def grad(x: Any, y: value[Any] | Sequence[value[Any]] | vector[Any] | tensor[Any
|
|||
elif opn == 'sqrt':
|
||||
add_grad(a, g * (0.5 / cp.sqrt(a)))
|
||||
|
||||
elif opn == 'abs':
|
||||
add_grad(a, g * cp.sign(a))
|
||||
|
||||
elif opn == 'neg':
|
||||
add_grad(a, -b)
|
||||
#elif opn == 'abs':
|
||||
# add_grad(x, g * cp.sign(x))
|
||||
|
||||
elif opn == 'sin':
|
||||
add_grad(a, g * cp.cos(a))
|
||||
|
|
@ -124,7 +121,7 @@ def grad(x: Any, y: value[Any] | Sequence[value[Any]] | vector[Any] | tensor[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.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 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 [grad_dict[yi.net] for yi in y]
|
||||
|
|
|
|||
|
|
@ -1,5 +1,5 @@
|
|||
import pkgutil
|
||||
from typing import Any, Sequence, TypeVar, overload, TypeAlias, Generic, Callable
|
||||
from typing import Any, Sequence, TypeVar, overload, TypeAlias, Generic, cast
|
||||
from ._stencils import stencil_database, detect_process_arch
|
||||
import copapy as cp
|
||||
from ._helper_types import TNum
|
||||
|
|
@ -230,11 +230,13 @@ class value(Generic[TNum]):
|
|||
def __rfloordiv__(self, other: NumLike) -> Any:
|
||||
return add_op('floordiv', [other, self])
|
||||
|
||||
def __abs__(self: 'value[TNum]') -> 'value[TNum]':
|
||||
return cp.abs(self)
|
||||
def __abs__(self: TCPNum) -> TCPNum:
|
||||
return cp.abs(self) # type: ignore
|
||||
|
||||
def __neg__(self: 'value[TNum]') -> 'value[TNum]':
|
||||
return add_op('neg', [self])
|
||||
def __neg__(self: TCPNum) -> TCPNum:
|
||||
if self.dtype == 'float':
|
||||
return cast(TCPNum, add_op('sub', [value(0.0), self]))
|
||||
return cast(TCPNum, add_op('sub', [value(0), self]))
|
||||
|
||||
def __gt__(self, other: TVarNumb) -> 'value[int]':
|
||||
return add_op('gt', [self, other], dtype='bool')
|
||||
|
|
@ -360,7 +362,7 @@ class CPConstant(Node):
|
|||
return self.node_hash
|
||||
|
||||
|
||||
class Store(Node):
|
||||
class Write(Node):
|
||||
def __init__(self, input: value[Any] | Net | int | float):
|
||||
if isinstance(input, value):
|
||||
net = input.net
|
||||
|
|
@ -370,7 +372,7 @@ class Store(Node):
|
|||
node = CPConstant(input)
|
||||
net = Net(node.dtype, node)
|
||||
|
||||
self.name = 'store_' + transl_type(net.dtype)
|
||||
self.name = 'write_' + transl_type(net.dtype)
|
||||
self.args = (net,)
|
||||
self.node_hash = hash(self.name) ^ hash(net.source.node_hash)
|
||||
|
||||
|
|
@ -429,15 +431,6 @@ class Op(Node):
|
|||
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
|
||||
|
|
|
|||
|
|
@ -2,7 +2,7 @@ from typing import Generator, Iterable, Any
|
|||
from . import _binwrite as binw
|
||||
from ._stencils import stencil_database, patch_entry
|
||||
from collections import defaultdict, deque
|
||||
from ._basic_types import Net, Node, Store, CPConstant, Op, transl_type
|
||||
from ._basic_types import Net, Node, Write, CPConstant, Op, transl_type
|
||||
|
||||
|
||||
def stable_toposort(edges: Iterable[tuple[Node, Node]]) -> list[Node]:
|
||||
|
|
@ -132,7 +132,7 @@ def get_const_nets(nodes: list[Node]) -> list[Net]:
|
|||
return [net_lookup[node] for node in nodes if isinstance(node, CPConstant)]
|
||||
|
||||
|
||||
def add_load_ops(node_list: list[Node]) -> Generator[tuple[Net | None, Node], None, None]:
|
||||
def add_read_ops(node_list: list[Node]) -> Generator[tuple[Net | None, Node], None, None]:
|
||||
"""Add read node before each op where arguments are not already positioned
|
||||
correctly in the registers
|
||||
|
||||
|
|
@ -156,7 +156,7 @@ def add_load_ops(node_list: list[Node]) -> Generator[tuple[Net | None, Node], No
|
|||
#if net in registers:
|
||||
# print('x swap registers')
|
||||
type_list = ['int' if r is None else transl_type(r.dtype) for r in registers]
|
||||
new_node = Op(f"load_{transl_type(net.dtype)}_reg{i}_" + '_'.join(type_list), [])
|
||||
new_node = Op(f"read_{transl_type(net.dtype)}_reg{i}_" + '_'.join(type_list), [])
|
||||
yield net, new_node
|
||||
registers[i] = net
|
||||
|
||||
|
|
@ -170,7 +170,7 @@ def add_load_ops(node_list: list[Node]) -> Generator[tuple[Net | None, Node], No
|
|||
yield None, node
|
||||
|
||||
|
||||
def add_store_ops(net_node_list: list[tuple[Net | None, Node]], const_nets: list[Net]) -> Generator[tuple[Net | None, Node], None, None]:
|
||||
def add_write_ops(net_node_list: list[tuple[Net | None, Node]], const_nets: list[Net]) -> Generator[tuple[Net | None, Node], None, None]:
|
||||
"""Add write operation for each new defined net if a read operation is later followed
|
||||
|
||||
Returns:
|
||||
|
|
@ -181,19 +181,19 @@ def add_store_ops(net_node_list: list[tuple[Net | None, Node]], const_nets: list
|
|||
# Initialize set of nets with constants
|
||||
stored_nets = set(const_nets)
|
||||
|
||||
#assert all(node.name.startswith('load_') for net, node in net_node_list if net)
|
||||
#assert all(node.name.startswith('read_') for net, node in net_node_list if net)
|
||||
read_back_nets = {
|
||||
net for net, node in net_node_list
|
||||
if net and node.name.startswith('load_')}
|
||||
if net and node.name.startswith('read_')}
|
||||
|
||||
registers: list[Net | None] = [None, None]
|
||||
|
||||
for net, node in net_node_list:
|
||||
if isinstance(node, Store):
|
||||
if isinstance(node, Write):
|
||||
assert len(registers) == 2
|
||||
type_list = [transl_type(r.dtype) if r else 'int' for r in registers]
|
||||
yield node.args[0], Op(f"store_{type_list[0]}_reg0_" + '_'.join(type_list), node.args)
|
||||
elif node.name.startswith('load_'):
|
||||
yield node.args[0], Op(f"write_{type_list[0]}_reg0_" + '_'.join(type_list), node.args)
|
||||
elif node.name.startswith('read_'):
|
||||
yield net, node
|
||||
else:
|
||||
yield None, node
|
||||
|
|
@ -207,7 +207,7 @@ def add_store_ops(net_node_list: list[tuple[Net | None, Node]], const_nets: list
|
|||
|
||||
if net in read_back_nets and net not in stored_nets:
|
||||
type_list = [transl_type(r.dtype) if r else 'int' for r in registers]
|
||||
yield net, Op(f"store_{type_list[0]}_reg0_" + '_'.join(type_list), [])
|
||||
yield net, Op(f"write_{type_list[0]}_reg0_" + '_'.join(type_list), [])
|
||||
stored_nets.add(net)
|
||||
|
||||
|
||||
|
|
@ -344,8 +344,8 @@ def compile_to_dag(node_list: Iterable[Node], sdb: stencil_database) -> tuple[bi
|
|||
|
||||
ordered_ops = list(stable_toposort(get_all_dag_edges(node_list)))
|
||||
const_net_list = get_const_nets(ordered_ops)
|
||||
output_ops = list(add_load_ops(ordered_ops))
|
||||
extended_output_ops = list(add_store_ops(output_ops, const_net_list))
|
||||
output_ops = list(add_read_ops(ordered_ops))
|
||||
extended_output_ops = list(add_write_ops(output_ops, const_net_list))
|
||||
|
||||
dw = binw.data_writer(sdb.byteorder)
|
||||
|
||||
|
|
|
|||
|
|
@ -1,7 +1,5 @@
|
|||
from . import vector
|
||||
from . import tensor
|
||||
from ._vectors import VecNumLike
|
||||
from ._tensors import TensorNumLike
|
||||
from . import value, NumLike
|
||||
from typing import TypeVar, Any, overload, Callable
|
||||
from ._basic_types import add_op, unifloat
|
||||
|
|
@ -17,8 +15,6 @@ def exp(x: float | int) -> float: ...
|
|||
def exp(x: value[Any]) -> value[float]: ...
|
||||
@overload
|
||||
def exp(x: vector[Any]) -> vector[float]: ...
|
||||
@overload
|
||||
def exp(x: tensor[Any]) -> tensor[float]: ...
|
||||
def exp(x: Any) -> Any:
|
||||
"""Exponential function to basis e
|
||||
|
||||
|
|
@ -30,7 +26,7 @@ def exp(x: Any) -> Any:
|
|||
"""
|
||||
if isinstance(x, value):
|
||||
return add_op('exp', [x])
|
||||
if isinstance(x, vector | tensor):
|
||||
if isinstance(x, vector):
|
||||
return x.map(exp)
|
||||
return float(math.exp(x))
|
||||
|
||||
|
|
@ -41,8 +37,6 @@ def log(x: float | int) -> float: ...
|
|||
def log(x: value[Any]) -> value[float]: ...
|
||||
@overload
|
||||
def log(x: vector[Any]) -> vector[float]: ...
|
||||
@overload
|
||||
def log(x: tensor[Any]) -> tensor[float]: ...
|
||||
def log(x: Any) -> Any:
|
||||
"""Logarithm to basis e
|
||||
|
||||
|
|
@ -54,7 +48,7 @@ def log(x: Any) -> Any:
|
|||
"""
|
||||
if isinstance(x, value):
|
||||
return add_op('log', [x])
|
||||
if isinstance(x, vector | tensor):
|
||||
if isinstance(x, vector):
|
||||
return x.map(log)
|
||||
return float(math.log(x))
|
||||
|
||||
|
|
@ -67,13 +61,7 @@ def pow(x: value[Any], y: NumLike) -> value[float]: ...
|
|||
def pow(x: NumLike, y: value[Any]) -> value[float]: ...
|
||||
@overload
|
||||
def pow(x: vector[Any], y: Any) -> vector[float]: ...
|
||||
@overload
|
||||
def pow(x: Any, y: vector[Any]) -> vector[float]: ...
|
||||
@overload
|
||||
def pow(x: tensor[Any], y: Any) -> tensor[float]: ...
|
||||
@overload
|
||||
def pow(x: Any, y: tensor[Any]) -> tensor[float]: ...
|
||||
def pow(x: TensorNumLike, y: TensorNumLike) -> Any:
|
||||
def pow(x: VecNumLike, y: VecNumLike) -> Any:
|
||||
"""x to the power of y
|
||||
|
||||
Arguments:
|
||||
|
|
@ -82,10 +70,8 @@ def pow(x: TensorNumLike, y: TensorNumLike) -> Any:
|
|||
Returns:
|
||||
result of x**y
|
||||
"""
|
||||
if isinstance(x, tensor) or isinstance(y, tensor):
|
||||
return _map2_tensor(x, y, pow)
|
||||
if isinstance(x, vector) or isinstance(y, vector):
|
||||
return _map2_vector(x, y, pow)
|
||||
return _map2(x, y, pow)
|
||||
if isinstance(y, int) and 0 <= y < 8:
|
||||
if y == 0:
|
||||
return 1
|
||||
|
|
@ -107,8 +93,6 @@ def sqrt(x: float | int) -> float: ...
|
|||
def sqrt(x: value[Any]) -> value[float]: ...
|
||||
@overload
|
||||
def sqrt(x: vector[Any]) -> vector[float]: ...
|
||||
@overload
|
||||
def sqrt(x: tensor[Any]) -> tensor[float]: ...
|
||||
def sqrt(x: Any) -> Any:
|
||||
"""Square root function
|
||||
|
||||
|
|
@ -120,7 +104,7 @@ def sqrt(x: Any) -> Any:
|
|||
"""
|
||||
if isinstance(x, value):
|
||||
return add_op('sqrt', [x])
|
||||
if isinstance(x, vector | tensor):
|
||||
if isinstance(x, vector):
|
||||
return x.map(sqrt)
|
||||
return float(math.sqrt(x))
|
||||
|
||||
|
|
@ -131,8 +115,6 @@ def sin(x: float | int) -> float: ...
|
|||
def sin(x: value[Any]) -> value[float]: ...
|
||||
@overload
|
||||
def sin(x: vector[Any]) -> vector[float]: ...
|
||||
@overload
|
||||
def sin(x: tensor[Any]) -> tensor[float]: ...
|
||||
def sin(x: Any) -> Any:
|
||||
"""Sine function
|
||||
|
||||
|
|
@ -144,7 +126,7 @@ def sin(x: Any) -> Any:
|
|||
"""
|
||||
if isinstance(x, value):
|
||||
return add_op('sin', [x])
|
||||
if isinstance(x, vector | tensor):
|
||||
if isinstance(x, vector):
|
||||
return x.map(sin)
|
||||
return math.sin(x)
|
||||
|
||||
|
|
@ -155,8 +137,6 @@ def cos(x: float | int) -> float: ...
|
|||
def cos(x: value[Any]) -> value[float]: ...
|
||||
@overload
|
||||
def cos(x: vector[Any]) -> vector[float]: ...
|
||||
@overload
|
||||
def cos(x: tensor[Any]) -> tensor[float]: ...
|
||||
def cos(x: Any) -> Any:
|
||||
"""Cosine function
|
||||
|
||||
|
|
@ -168,7 +148,7 @@ def cos(x: Any) -> Any:
|
|||
"""
|
||||
if isinstance(x, value):
|
||||
return add_op('cos', [x])
|
||||
if isinstance(x, vector | tensor):
|
||||
if isinstance(x, vector):
|
||||
return x.map(cos)
|
||||
return math.cos(x)
|
||||
|
||||
|
|
@ -179,8 +159,6 @@ def tan(x: float | int) -> float: ...
|
|||
def tan(x: value[Any]) -> value[float]: ...
|
||||
@overload
|
||||
def tan(x: vector[Any]) -> vector[float]: ...
|
||||
@overload
|
||||
def tan(x: tensor[Any]) -> tensor[float]: ...
|
||||
def tan(x: Any) -> Any:
|
||||
"""Tangent function
|
||||
|
||||
|
|
@ -192,7 +170,8 @@ def tan(x: Any) -> Any:
|
|||
"""
|
||||
if isinstance(x, value):
|
||||
return add_op('tan', [x])
|
||||
if isinstance(x, vector | tensor):
|
||||
if isinstance(x, vector):
|
||||
#return x.map(tan)
|
||||
return x.map(tan)
|
||||
return math.tan(x)
|
||||
|
||||
|
|
@ -203,8 +182,6 @@ def atan(x: float | int) -> float: ...
|
|||
def atan(x: value[Any]) -> value[float]: ...
|
||||
@overload
|
||||
def atan(x: vector[Any]) -> vector[float]: ...
|
||||
@overload
|
||||
def atan(x: tensor[Any]) -> tensor[float]: ...
|
||||
def atan(x: Any) -> Any:
|
||||
"""Inverse tangent function
|
||||
|
||||
|
|
@ -216,7 +193,7 @@ def atan(x: Any) -> Any:
|
|||
"""
|
||||
if isinstance(x, value):
|
||||
return add_op('atan', [x])
|
||||
if isinstance(x, vector | tensor):
|
||||
if isinstance(x, vector):
|
||||
return x.map(atan)
|
||||
return math.atan(x)
|
||||
|
||||
|
|
@ -231,11 +208,7 @@ def atan2(x: NumLike, y: value[Any]) -> value[float]: ...
|
|||
def atan2(x: vector[float], y: VecNumLike) -> vector[float]: ...
|
||||
@overload
|
||||
def atan2(x: VecNumLike, y: vector[float]) -> vector[float]: ...
|
||||
@overload
|
||||
def atan2(x: tensor[float], y: TensorNumLike) -> tensor[float]: ...
|
||||
@overload
|
||||
def atan2(x: TensorNumLike, y: tensor[float]) -> tensor[float]: ...
|
||||
def atan2(x: TensorNumLike, y: TensorNumLike) -> Any:
|
||||
def atan2(x: VecNumLike, y: VecNumLike) -> Any:
|
||||
"""2-argument arctangent
|
||||
|
||||
Arguments:
|
||||
|
|
@ -245,10 +218,8 @@ def atan2(x: TensorNumLike, y: TensorNumLike) -> Any:
|
|||
Returns:
|
||||
Result in radian
|
||||
"""
|
||||
if isinstance(x, tensor) or isinstance(y, tensor):
|
||||
return _map2_tensor(x, y, atan2)
|
||||
if isinstance(x, vector) or isinstance(y, vector):
|
||||
return _map2_vector(x, y, atan2)
|
||||
return _map2(x, y, atan2)
|
||||
if isinstance(x, value) or isinstance(y, value):
|
||||
return add_op('atan2', [x, y])
|
||||
return math.atan2(x, y)
|
||||
|
|
@ -260,8 +231,6 @@ def asin(x: float | int) -> float: ...
|
|||
def asin(x: value[Any]) -> value[float]: ...
|
||||
@overload
|
||||
def asin(x: vector[Any]) -> vector[float]: ...
|
||||
@overload
|
||||
def asin(x: tensor[Any]) -> tensor[float]: ...
|
||||
def asin(x: Any) -> Any:
|
||||
"""Inverse sine function
|
||||
|
||||
|
|
@ -273,7 +242,7 @@ def asin(x: Any) -> Any:
|
|||
"""
|
||||
if isinstance(x, value):
|
||||
return add_op('asin', [x])
|
||||
if isinstance(x, vector | tensor):
|
||||
if isinstance(x, vector):
|
||||
return x.map(asin)
|
||||
return math.asin(x)
|
||||
|
||||
|
|
@ -284,8 +253,6 @@ def acos(x: float | int) -> float: ...
|
|||
def acos(x: value[Any]) -> value[float]: ...
|
||||
@overload
|
||||
def acos(x: vector[Any]) -> vector[float]: ...
|
||||
@overload
|
||||
def acos(x: tensor[Any]) -> tensor[float]: ...
|
||||
def acos(x: Any) -> Any:
|
||||
"""Inverse cosine function
|
||||
|
||||
|
|
@ -297,12 +264,11 @@ def acos(x: Any) -> Any:
|
|||
"""
|
||||
if isinstance(x, value):
|
||||
return add_op('acos', [x])
|
||||
if isinstance(x, vector | tensor):
|
||||
if isinstance(x, vector):
|
||||
return x.map(acos)
|
||||
return math.asin(x)
|
||||
|
||||
|
||||
# Debug test function
|
||||
@overload
|
||||
def get_42(x: float | int) -> float: ...
|
||||
@overload
|
||||
|
|
@ -320,9 +286,7 @@ def abs(x: U) -> U: ...
|
|||
def abs(x: value[U]) -> value[U]: ...
|
||||
@overload
|
||||
def abs(x: vector[U]) -> vector[U]: ...
|
||||
@overload
|
||||
def abs(x: tensor[U]) -> tensor[U]: ...
|
||||
def abs(x: U | value[U] | vector[U] | tensor[U]) -> Any:
|
||||
def abs(x: U | value[U] | vector[U]) -> Any:
|
||||
"""Absolute value function
|
||||
|
||||
Arguments:
|
||||
|
|
@ -333,20 +297,18 @@ def abs(x: U | value[U] | vector[U] | tensor[U]) -> Any:
|
|||
"""
|
||||
if isinstance(x, value):
|
||||
return add_op('abs', [x])
|
||||
if isinstance(x, vector | tensor):
|
||||
if isinstance(x, vector):
|
||||
return x.map(abs)
|
||||
return (x < 0) * -x + (x >= 0) * x
|
||||
|
||||
|
||||
@overload
|
||||
def sign(x: U) -> int: ...
|
||||
def sign(x: U) -> U: ...
|
||||
@overload
|
||||
def sign(x: value[U]) -> value[int]: ...
|
||||
def sign(x: value[U]) -> value[U]: ...
|
||||
@overload
|
||||
def sign(x: vector[U]) -> vector[int]: ...
|
||||
@overload
|
||||
def sign(x: tensor[U]) -> tensor[int]: ...
|
||||
def sign(x: U | value[U] | vector[U] | tensor[U]) -> Any:
|
||||
def sign(x: vector[U]) -> vector[U]: ...
|
||||
def sign(x: U | value[U] | vector[U]) -> Any:
|
||||
"""Return 1 for positive numbers and -1 for negative numbers.
|
||||
For an input of 0 the return value is 0.
|
||||
|
||||
|
|
@ -356,11 +318,8 @@ def sign(x: U | value[U] | vector[U] | tensor[U]) -> Any:
|
|||
Returns:
|
||||
-1, 0 or 1
|
||||
"""
|
||||
if isinstance(x, value):
|
||||
return add_op('sign', [x])
|
||||
if isinstance(x, vector | tensor):
|
||||
return x.map(sign)
|
||||
return (x > 0) - (x < 0)
|
||||
ret = (x > 0) - (x < 0)
|
||||
return ret
|
||||
|
||||
|
||||
@overload
|
||||
|
|
@ -408,13 +367,7 @@ def min(x: U | value[U], y: U | value[U]) -> Any:
|
|||
Returns:
|
||||
Minimum of x and y
|
||||
"""
|
||||
if isinstance(x, value):
|
||||
return add_op('min', [x, y])
|
||||
if isinstance(x, tensor):
|
||||
return _map2_tensor(x, y, min)
|
||||
if isinstance(x, vector):
|
||||
return _map2_vector(x, y, min)
|
||||
return x if x < y else y
|
||||
return (x < y) * x + (x >= y) * y
|
||||
|
||||
|
||||
@overload
|
||||
|
|
@ -433,13 +386,7 @@ def max(x: U | value[U], y: U | value[U]) -> Any:
|
|||
Returns:
|
||||
Maximum of x and y
|
||||
"""
|
||||
if isinstance(x, value):
|
||||
return add_op('max', [x, y])
|
||||
if isinstance(x, tensor):
|
||||
return _map2_tensor(x, y, max)
|
||||
if isinstance(x, vector):
|
||||
return _map2_vector(x, y, max)
|
||||
return x if x > y else y
|
||||
return (x > y) * x + (x <= y) * y
|
||||
|
||||
|
||||
@overload
|
||||
|
|
@ -453,16 +400,7 @@ def lerp(v1: U, v2: U, t: float) -> U: ...
|
|||
@overload
|
||||
def lerp(v1: vector[U], v2: vector[U], t: unifloat) -> vector[U]: ...
|
||||
def lerp(v1: U | value[U] | vector[U], v2: U | value[U] | vector[U], t: unifloat) -> Any:
|
||||
"""Linearly interpolate between two values or vectors v1 and v2 by a factor t.
|
||||
|
||||
Arguments:
|
||||
v1: First value or vector
|
||||
v2: Second value or vector
|
||||
t: Interpolation factor (0.0 to 1.0)
|
||||
|
||||
Returns:
|
||||
Interpolated value or vector
|
||||
"""
|
||||
"""Linearly interpolate between two values or vectors v1 and v2 by a factor t."""
|
||||
if isinstance(v1, vector) or isinstance(v2, vector):
|
||||
assert isinstance(v1, vector) and isinstance(v2, vector), "None or both v1 and v2 must be vectors."
|
||||
assert len(v1.values) == len(v2.values), "Vectors must be of the same length."
|
||||
|
|
@ -476,15 +414,13 @@ def relu(x: U) -> U: ...
|
|||
def relu(x: value[U]) -> value[U]: ...
|
||||
@overload
|
||||
def relu(x: vector[U]) -> vector[U]: ...
|
||||
@overload
|
||||
def relu(x: tensor[U]) -> tensor[U]: ...
|
||||
def relu(x: U | value[U] | vector[U] | tensor[U]) -> Any:
|
||||
def relu(x: U | value[U] | vector[U]) -> Any:
|
||||
"""Returns x for x > 0 and otherwise 0."""
|
||||
ret = x * (x > 0)
|
||||
ret = (x > 0) * x
|
||||
return ret
|
||||
|
||||
|
||||
def _map2_vector(self: VecNumLike, other: VecNumLike, func: Callable[[Any, Any], value[U] | U]) -> vector[U]:
|
||||
def _map2(self: VecNumLike, other: VecNumLike, func: Callable[[Any, Any], value[U] | U]) -> vector[U]:
|
||||
"""Applies a function to each element of the vector and a second vector or scalar."""
|
||||
if isinstance(self, vector) and isinstance(other, vector):
|
||||
return vector(func(x, y) for x, y in zip(self.values, other.values))
|
||||
|
|
@ -494,20 +430,3 @@ def _map2_vector(self: VecNumLike, other: VecNumLike, func: Callable[[Any, Any],
|
|||
return vector(func(self, x) for x in other.values)
|
||||
else:
|
||||
return vector([func(self, other)])
|
||||
|
||||
|
||||
def _map2_tensor(self: TensorNumLike, other: TensorNumLike, func: Callable[[Any, Any], value[U] | U]) -> tensor[U]:
|
||||
"""Applies a function to each element of the vector and a second vector or scalar."""
|
||||
if isinstance(self, vector):
|
||||
self = tensor(self.values, (len(self.values),))
|
||||
if isinstance(other, vector):
|
||||
other = tensor(other.values, (len(other.values),))
|
||||
if isinstance(self, tensor) and isinstance(other, tensor):
|
||||
assert self.shape == other.shape, "Tensors must have the same shape"
|
||||
return tensor([func(x, y) for x, y in zip(self.values, other.values)], self.shape)
|
||||
elif isinstance(self, tensor):
|
||||
return tensor([func(x, other) for x in self.values], self.shape)
|
||||
elif isinstance(other, tensor):
|
||||
return tensor([func(self, x) for x in other.values], other.shape)
|
||||
else:
|
||||
return tensor(func(self, other))
|
||||
|
|
|
|||
|
|
@ -0,0 +1,359 @@
|
|||
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,7 +2,8 @@ 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 value, Net, Node, Store, NumLike, ArrayType, stencil_db_from_package
|
||||
from ._basic_types import stencil_db_from_package
|
||||
from ._basic_types import value, Net, Node, Write, NumLike
|
||||
from ._compiler import compile_to_dag
|
||||
|
||||
T = TypeVar("T", int, float)
|
||||
|
|
@ -65,7 +66,7 @@ class Target():
|
|||
def __del__(self) -> None:
|
||||
clear_target(self._context)
|
||||
|
||||
def compile(self, *values: NumLike | value[T] | ArrayType[T] | Iterable[T | value[T]]) -> None:
|
||||
def compile(self, *values: int | float | value[Any] | Iterable[int | float | value[Any]]) -> None:
|
||||
"""Compiles the code to compute the given values.
|
||||
|
||||
Arguments:
|
||||
|
|
@ -73,16 +74,13 @@ class Target():
|
|||
"""
|
||||
nodes: list[Node] = []
|
||||
for input in values:
|
||||
if isinstance(input, ArrayType):
|
||||
for v in input.values:
|
||||
if isinstance(v, value):
|
||||
nodes.append(Store(v))
|
||||
elif isinstance(input, Iterable):
|
||||
if isinstance(input, Iterable):
|
||||
for v in input:
|
||||
if isinstance(v, value):
|
||||
nodes.append(Store(v))
|
||||
elif isinstance(input, value):
|
||||
nodes.append(Store(input))
|
||||
nodes.append(Write(v))
|
||||
else:
|
||||
if isinstance(input, value):
|
||||
nodes.append(Write(input))
|
||||
|
||||
dw, self._values = compile_to_dag(nodes, self.sdb)
|
||||
dw.write_com(binw.Command.END_COM)
|
||||
|
|
@ -102,9 +100,7 @@ class Target():
|
|||
def read_value(self, variables: NumLike) -> float | int | bool: ...
|
||||
@overload
|
||||
def read_value(self, variables: Iterable[T | value[T]]) -> list[T]: ...
|
||||
@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:
|
||||
def read_value(self, variables: NumLike | value[T] | Iterable[T | value[T]]) -> Any:
|
||||
"""Reads the numeric value of a copapy type.
|
||||
|
||||
Arguments:
|
||||
|
|
@ -113,9 +109,6 @@ 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]
|
||||
|
||||
|
|
|
|||
|
|
@ -1,955 +0,0 @@
|
|||
from copapy._basic_types import NumLike, ArrayType
|
||||
from . import value
|
||||
from ._vectors import vector, VecFloatLike, VecIntLike, VecNumLike
|
||||
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] | vector[int] | value[int] | int | bool'
|
||||
TensorFloatLike: TypeAlias = 'tensor[float] | vector[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.
|
||||
shape: Optional shape of the tensor. If not provided, inferred from values.
|
||||
"""
|
||||
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[Any]': ...
|
||||
def __add__(self, other: TensorNumLike) -> Any:
|
||||
"""Element-wise addition."""
|
||||
return self._binary_op(other, lambda a, b: a + b)
|
||||
|
||||
@overload
|
||||
def __radd__(self: 'tensor[int]', other: VecFloatLike) -> 'tensor[float]': ...
|
||||
@overload
|
||||
def __radd__(self: 'tensor[int]', other: VecIntLike) -> 'tensor[int]': ...
|
||||
@overload
|
||||
def __radd__(self: 'tensor[float]', other: VecNumLike) -> 'tensor[float]': ...
|
||||
@overload
|
||||
def __radd__(self, other: VecNumLike) -> 'tensor[Any]': ...
|
||||
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[Any]': ...
|
||||
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[int]', other: VecFloatLike) -> 'tensor[float]': ...
|
||||
@overload
|
||||
def __rsub__(self: 'tensor[int]', other: VecIntLike) -> 'tensor[int]': ...
|
||||
@overload
|
||||
def __rsub__(self: 'tensor[float]', other: VecNumLike) -> 'tensor[float]': ...
|
||||
@overload
|
||||
def __rsub__(self, other: VecNumLike) -> 'tensor[Any]': ...
|
||||
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[Any]': ...
|
||||
def __mul__(self, other: TensorNumLike) -> Any:
|
||||
"""Element-wise multiplication."""
|
||||
return self._binary_op(other, lambda a, b: a * b)
|
||||
|
||||
@overload
|
||||
def __rmul__(self: 'tensor[int]', other: VecFloatLike) -> 'tensor[float]': ...
|
||||
@overload
|
||||
def __rmul__(self: 'tensor[int]', other: VecIntLike) -> 'tensor[int]': ...
|
||||
@overload
|
||||
def __rmul__(self: 'tensor[float]', other: VecNumLike) -> 'tensor[float]': ...
|
||||
@overload
|
||||
def __rmul__(self, other: VecNumLike) -> 'tensor[Any]': ...
|
||||
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[Any]': ...
|
||||
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[int]', other: VecFloatLike) -> 'tensor[float]': ...
|
||||
@overload
|
||||
def __rpow__(self: 'tensor[int]', other: VecIntLike) -> 'tensor[int]': ...
|
||||
@overload
|
||||
def __rpow__(self: 'tensor[float]', other: VecNumLike) -> 'tensor[float]': ...
|
||||
@overload
|
||||
def __rpow__(self, other: VecNumLike) -> 'tensor[Any]': ...
|
||||
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,9 +1,8 @@
|
|||
from . import value
|
||||
from ._mixed import mixed_sum, mixed_homogenize
|
||||
from typing import Sequence, TypeVar, Iterable, Any, overload, TypeAlias, Callable, Iterator
|
||||
from typing import Sequence, TypeVar, Iterable, Any, overload, TypeAlias, Callable, Iterator, Generic
|
||||
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'
|
||||
|
|
@ -14,13 +13,8 @@ U = TypeVar("U", int, float)
|
|||
epsilon = 1e-20
|
||||
|
||||
|
||||
class vector(ArrayType[TNum]):
|
||||
class vector(Generic[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.
|
||||
|
|
@ -29,8 +23,6 @@ class vector(ArrayType[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})"
|
||||
|
|
@ -53,10 +45,6 @@ class vector(ArrayType[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
|
||||
|
|
@ -280,6 +268,10 @@ class vector(ArrayType[TNum]):
|
|||
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]: ...
|
||||
|
|
|
|||
|
|
@ -4,10 +4,10 @@ and give access to compiler internals and debugging tools.
|
|||
"""
|
||||
|
||||
from ._target import add_read_command
|
||||
from ._basic_types import Net, Op, Node, CPConstant, Store, stencil_db_from_package
|
||||
from ._basic_types import Net, Op, Node, CPConstant, Write, stencil_db_from_package
|
||||
from ._compiler import compile_to_dag, \
|
||||
stable_toposort, get_const_nets, get_all_dag_edges, add_load_ops, get_all_dag_edges_between, \
|
||||
add_store_ops, get_dag_stats
|
||||
stable_toposort, get_const_nets, get_all_dag_edges, add_read_ops, get_all_dag_edges_between, \
|
||||
add_write_ops, get_dag_stats
|
||||
|
||||
__all__ = [
|
||||
"add_read_command",
|
||||
|
|
@ -15,14 +15,14 @@ __all__ = [
|
|||
"Op",
|
||||
"Node",
|
||||
"CPConstant",
|
||||
"Store",
|
||||
"Write",
|
||||
"compile_to_dag",
|
||||
"stable_toposort",
|
||||
"get_const_nets",
|
||||
"get_all_dag_edges",
|
||||
"get_all_dag_edges_between",
|
||||
"add_load_ops",
|
||||
"add_store_ops",
|
||||
"add_read_ops",
|
||||
"add_write_ops",
|
||||
"stencil_db_from_package",
|
||||
"get_dag_stats"
|
||||
]
|
||||
|
|
|
|||
|
|
@ -83,15 +83,6 @@ 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"""
|
||||
|
|
@ -166,42 +157,8 @@ def get_floordiv(op: str, type1: str, type2: str) -> str:
|
|||
"""
|
||||
else:
|
||||
return f"""
|
||||
STENCIL void {op}_{type1}_{type2}({type1} a, {type2} b) {{
|
||||
result_float_{type2}(floorf((float)a / (float)b), b);
|
||||
}}
|
||||
"""
|
||||
|
||||
|
||||
@norm_indent
|
||||
def get_min(type1: str, type2: str) -> str:
|
||||
if type1 == 'int' and type2 == 'int':
|
||||
return f"""
|
||||
STENCIL void min_{type1}_{type2}({type1} a, {type2} b) {{
|
||||
result_int_{type2}(a < b ? a : b, b);
|
||||
}}
|
||||
"""
|
||||
else:
|
||||
return f"""
|
||||
STENCIL void min_{type1}_{type2}({type1} a, {type2} b) {{
|
||||
float _a = (float)a; float _b = (float)b;
|
||||
result_float_{type2}(_a < _b ? _a : _b, b);
|
||||
}}
|
||||
"""
|
||||
|
||||
|
||||
@norm_indent
|
||||
def get_max(type1: str, type2: str) -> str:
|
||||
if type1 == 'int' and type2 == 'int':
|
||||
return f"""
|
||||
STENCIL void max_{type1}_{type2}({type1} a, {type2} b) {{
|
||||
result_int_{type2}(a > b ? a : b, b);
|
||||
}}
|
||||
"""
|
||||
else:
|
||||
return f"""
|
||||
STENCIL void max_{type1}_{type2}({type1} a, {type2} b) {{
|
||||
float _a = (float)a; float _b = (float)b;
|
||||
result_float_{type2}(_a > _b ? _a : _b, b);
|
||||
STENCIL void {op}_{type1}_{type2}({type1} arg1, {type2} arg2) {{
|
||||
result_float_{type2}(floorf((float)arg1 / (float)arg2), arg2);
|
||||
}}
|
||||
"""
|
||||
|
||||
|
|
@ -221,27 +178,27 @@ def get_result_stubs2(type1: str, type2: str) -> str:
|
|||
|
||||
|
||||
@norm_indent
|
||||
def get_load_reg0_code(type1: str, type2: str, type_out: str) -> str:
|
||||
def get_read_reg0_code(type1: str, type2: str, type_out: str) -> str:
|
||||
return f"""
|
||||
STENCIL void load_{type_out}_reg0_{type1}_{type2}({type1} arg1, {type2} arg2) {{
|
||||
STENCIL void read_{type_out}_reg0_{type1}_{type2}({type1} arg1, {type2} arg2) {{
|
||||
result_{type_out}_{type2}(dummy_{type_out}, arg2);
|
||||
}}
|
||||
"""
|
||||
|
||||
|
||||
@norm_indent
|
||||
def get_load_reg1_code(type1: str, type2: str, type_out: str) -> str:
|
||||
def get_read_reg1_code(type1: str, type2: str, type_out: str) -> str:
|
||||
return f"""
|
||||
STENCIL void load_{type_out}_reg1_{type1}_{type2}({type1} arg1, {type2} arg2) {{
|
||||
STENCIL void read_{type_out}_reg1_{type1}_{type2}({type1} arg1, {type2} arg2) {{
|
||||
result_{type1}_{type_out}(arg1, dummy_{type_out});
|
||||
}}
|
||||
"""
|
||||
|
||||
|
||||
@norm_indent
|
||||
def get_store_code(type1: str, type2: str) -> str:
|
||||
def get_write_code(type1: str, type2: str) -> str:
|
||||
return f"""
|
||||
STENCIL void store_{type1}_reg0_{type1}_{type2}({type1} arg1, {type2} arg2) {{
|
||||
STENCIL void write_{type1}_reg0_{type1}_{type2}({type1} arg1, {type2} arg2) {{
|
||||
dummy_{type1} = arg1;
|
||||
result_{type1}_{type2}(arg1, arg2);
|
||||
}}
|
||||
|
|
@ -292,9 +249,6 @@ 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)
|
||||
|
|
@ -302,17 +256,10 @@ if __name__ == "__main__":
|
|||
code += get_math_func1('fabsf', 'float', 'abs')
|
||||
code += get_custom_stencil('abs_int(int arg1)', 'result_int(__builtin_abs(arg1));')
|
||||
|
||||
for t in types:
|
||||
code += get_custom_stencil(f"sign_{t}({t} arg1)", f"result_int((arg1 > 0) - (arg1 < 0));")
|
||||
|
||||
fnames = ['atan2', 'pow']
|
||||
for fn, t1, t2 in permutate(fnames, types, types):
|
||||
code += get_math_func2(fn, t1, t2)
|
||||
|
||||
for t1, t2 in permutate(types, types):
|
||||
code += get_min(t1, t2)
|
||||
code += get_max(t1, t2)
|
||||
|
||||
for op, t1, t2 in permutate(ops, types, types):
|
||||
t_out = t1 if t1 == t2 else 'float'
|
||||
if op == 'floordiv':
|
||||
|
|
@ -330,11 +277,11 @@ if __name__ == "__main__":
|
|||
code += get_op_code('mod', 'int', 'int', 'int')
|
||||
|
||||
for t1, t2, t_out in permutate(types, types, types):
|
||||
code += get_load_reg0_code(t1, t2, t_out)
|
||||
code += get_load_reg1_code(t1, t2, t_out)
|
||||
code += get_read_reg0_code(t1, t2, t_out)
|
||||
code += get_read_reg1_code(t1, t2, t_out)
|
||||
|
||||
for t1, t2 in permutate(types, types):
|
||||
code += get_store_code(t1, t2)
|
||||
code += get_write_code(t1, t2)
|
||||
|
||||
print(f"Write file {args.path}...")
|
||||
with open(args.path, 'w') as f:
|
||||
|
|
|
|||
|
|
@ -1,5 +1,5 @@
|
|||
from copapy import value
|
||||
from copapy.backend import Store
|
||||
from copapy.backend import Write
|
||||
import copapy.backend as cpb
|
||||
|
||||
|
||||
|
|
@ -19,16 +19,16 @@ def test_ast_generation():
|
|||
#i1 = c1 * 2
|
||||
#r1 = i1 + 7
|
||||
#r2 = i1 + 9
|
||||
#out = [Store(r1), Store(r2)]
|
||||
#out = [Write(r1), Write(r2)]
|
||||
|
||||
c1 = value(4)
|
||||
c2 = value(2)
|
||||
#i1 = c1 * 2
|
||||
#r1 = i1 + 7 + (c2 + 7 * 9)
|
||||
#r2 = i1 + 9
|
||||
#out = [Store(r1), Store(r2)]
|
||||
#out = [Write(r1), Write(r2)]
|
||||
r1 = c1 * 5 + 8 + c2 * 3
|
||||
out = [Store(r1)]
|
||||
out = [Write(r1)]
|
||||
|
||||
print(out)
|
||||
print('-- get_edges:')
|
||||
|
|
@ -48,12 +48,12 @@ def test_ast_generation():
|
|||
print('#', p)
|
||||
|
||||
print('-- add_read_ops:')
|
||||
output_ops = list(cpb.add_load_ops(ordered_ops))
|
||||
output_ops = list(cpb.add_read_ops(ordered_ops))
|
||||
for p in output_ops:
|
||||
print('#', p)
|
||||
|
||||
print('-- add_write_ops:')
|
||||
extended_output_ops = list(cpb.add_store_ops(output_ops, const_list))
|
||||
extended_output_ops = list(cpb.add_write_ops(output_ops, const_list))
|
||||
for p in extended_output_ops:
|
||||
print('#', p)
|
||||
print('--')
|
||||
|
|
|
|||
|
|
@ -13,7 +13,7 @@ def test_autograd():
|
|||
c += c + 1
|
||||
c += 1 + c + (-a)
|
||||
d += d * 2 + cp.relu(b + a)
|
||||
d += 3 * d + cp.relu(-a + b)
|
||||
d += 3 * d + cp.relu(b - a)
|
||||
e = c - d
|
||||
f = e**2
|
||||
g = f / 2.0
|
||||
|
|
@ -34,26 +34,5 @@ def test_autograd():
|
|||
assert pytest.approx(dg[1], abs=1e-4) == 645.57725 # pyright: ignore[reportUnknownMemberType]
|
||||
|
||||
|
||||
def test_autograd_extended():
|
||||
a = value(-4.0)
|
||||
b = value(2.0)
|
||||
c = a + b
|
||||
d = a * b + b**3
|
||||
c += c + 1
|
||||
c += 1 + c + (-a)
|
||||
d += d * 2 + cp.relu(b + a)
|
||||
d += 3 * d + cp.relu(b - a)
|
||||
e = c - cp.sin(-d)
|
||||
f = cp.abs(e**2)
|
||||
g = f / 2.0
|
||||
g += 10.0 / f
|
||||
|
||||
dg = grad(g, (a, b))
|
||||
|
||||
tg = cp.Target()
|
||||
tg.compile(g, dg)
|
||||
tg.run()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
test_autograd()
|
||||
|
|
|
|||
|
|
@ -1,5 +1,5 @@
|
|||
from copapy import value
|
||||
from copapy.backend import Store, compile_to_dag, add_read_command
|
||||
from copapy.backend import Write, compile_to_dag, add_read_command
|
||||
import copapy as cp
|
||||
import subprocess
|
||||
from copapy import _binwrite
|
||||
|
|
@ -22,7 +22,7 @@ def test_compile():
|
|||
# Function with no passing-on-jump as last instruction:
|
||||
ret_test = [r for v in test_vals for r in (cp.tan(value(v)),)]
|
||||
|
||||
out = [Store(r) for r in ret_test]
|
||||
out = [Write(r) for r in ret_test]
|
||||
|
||||
il, variables = compile_to_dag(out, copapy.generic_sdb)
|
||||
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
import time
|
||||
from copapy import backend
|
||||
from copapy.backend import Store, stencil_db_from_package
|
||||
from copapy.backend import Write, stencil_db_from_package
|
||||
import copapy.backend as cpb
|
||||
import copapy as cp
|
||||
import copapy._binwrite as binw
|
||||
|
|
@ -13,7 +13,7 @@ def test_timing_compiler():
|
|||
#t2 = t1.sum()
|
||||
t3 = cp.vector(cp.value(1 / (v + 1)) for v in range(256))
|
||||
t5 = ((t3 * t1) * 2).magnitude()
|
||||
out = [Store(t5)]
|
||||
out = [Write(t5)]
|
||||
|
||||
print(out)
|
||||
|
||||
|
|
@ -45,7 +45,7 @@ def test_timing_compiler():
|
|||
|
||||
print('-- add_read_ops:')
|
||||
t0 = time.time()
|
||||
output_ops = list(cpb.add_load_ops(ordered_ops))
|
||||
output_ops = list(cpb.add_read_ops(ordered_ops))
|
||||
t1 = time.time()
|
||||
#for p in output_ops:
|
||||
# print('#', p)
|
||||
|
|
@ -53,7 +53,7 @@ def test_timing_compiler():
|
|||
|
||||
print('-- add_write_ops:')
|
||||
t0 = time.time()
|
||||
extended_output_ops = list(cpb.add_store_ops(output_ops, const_net_list))
|
||||
extended_output_ops = list(cpb.add_write_ops(output_ops, const_net_list))
|
||||
t1 = time.time()
|
||||
#for p in extended_output_ops:
|
||||
# print('#', p)
|
||||
|
|
|
|||
|
|
@ -1,5 +1,5 @@
|
|||
from copapy import NumLike
|
||||
from copapy.backend import Store, compile_to_dag, add_read_command
|
||||
from copapy.backend import Write, compile_to_dag, add_read_command
|
||||
import copapy as cp
|
||||
import subprocess
|
||||
import struct
|
||||
|
|
@ -58,7 +58,7 @@ def test_compile():
|
|||
|
||||
ret = (t2, t4, t5)
|
||||
|
||||
out = [Store(r) for r in ret]
|
||||
out = [Write(r) for r in ret]
|
||||
|
||||
il, variables = compile_to_dag(out, copapy.generic_sdb)
|
||||
|
||||
|
|
|
|||
|
|
@ -1,5 +1,5 @@
|
|||
from copapy import NumLike
|
||||
from copapy.backend import Store, compile_to_dag, add_read_command
|
||||
from copapy.backend import Write, compile_to_dag, add_read_command
|
||||
import subprocess
|
||||
from copapy import _binwrite
|
||||
import copapy.backend as backend
|
||||
|
|
@ -52,7 +52,7 @@ def test_compile():
|
|||
|
||||
ret = (t2, t4, t5)
|
||||
|
||||
out = [Store(r) for r in ret]
|
||||
out = [Write(r) for r in ret]
|
||||
|
||||
sdb = backend.stencil_db_from_package('arm64')
|
||||
il, variables = compile_to_dag(out, sdb)
|
||||
|
|
|
|||
|
|
@ -1,5 +1,5 @@
|
|||
from copapy import NumLike
|
||||
from copapy.backend import Store, compile_to_dag, add_read_command
|
||||
from copapy.backend import Write, compile_to_dag, add_read_command
|
||||
import subprocess
|
||||
from copapy import _binwrite
|
||||
import copapy.backend as backend
|
||||
|
|
@ -52,7 +52,7 @@ def test_compile():
|
|||
|
||||
ret = (t2, t4, t5)
|
||||
|
||||
out = [Store(r) for r in ret]
|
||||
out = [Write(r) for r in ret]
|
||||
|
||||
sdb = backend.stencil_db_from_package('armv7')
|
||||
il, variables = compile_to_dag(out, sdb)
|
||||
|
|
|
|||
|
|
@ -1,5 +1,5 @@
|
|||
from copapy import value, NumLike
|
||||
from copapy.backend import Store, compile_to_dag, add_read_command
|
||||
from copapy.backend import Write, compile_to_dag, add_read_command, Net
|
||||
import copapy
|
||||
import subprocess
|
||||
from copapy import _binwrite
|
||||
|
|
@ -26,7 +26,7 @@ def test_compile():
|
|||
|
||||
ret = function(c1)
|
||||
|
||||
out = [Store(r) for r in ret]
|
||||
out = [Write(r) for r in ret]
|
||||
|
||||
il, vars = compile_to_dag(out, copapy.generic_sdb)
|
||||
|
||||
|
|
|
|||
|
|
@ -1,5 +1,5 @@
|
|||
from copapy import value
|
||||
from copapy.backend import Store, compile_to_dag, add_read_command
|
||||
from copapy.backend import Write, compile_to_dag, add_read_command
|
||||
import copapy as cp
|
||||
import subprocess
|
||||
from copapy import _binwrite
|
||||
|
|
@ -21,7 +21,7 @@ def test_compile_sqrt():
|
|||
ret = [r for v in test_vals for r in (cp.sqrt(value(v)),)]
|
||||
|
||||
|
||||
out = [Store(r) for r in ret]
|
||||
out = [Write(r) for r in ret]
|
||||
|
||||
il, variables = compile_to_dag(out, copapy.generic_sdb)
|
||||
|
||||
|
|
@ -55,7 +55,7 @@ def test_compile_log():
|
|||
ret = [r for v in test_vals for r in (cp.log(value(v)),)]
|
||||
|
||||
|
||||
out = [Store(r) for r in ret]
|
||||
out = [Write(r) for r in ret]
|
||||
|
||||
il, variables = compile_to_dag(out, copapy.generic_sdb)
|
||||
|
||||
|
|
@ -89,7 +89,7 @@ def test_compile_sin():
|
|||
ret = [r for v in test_vals for r in (cp.sin(value(v)),)]
|
||||
|
||||
|
||||
out = [Store(r) for r in ret]
|
||||
out = [Write(r) for r in ret]
|
||||
|
||||
il, variables = compile_to_dag(out, copapy.generic_sdb)
|
||||
|
||||
|
|
|
|||
|
|
@ -1,12 +1,12 @@
|
|||
import copapy as cp
|
||||
from copapy import value
|
||||
from copapy.backend import get_dag_stats, Store
|
||||
from copapy.backend import get_dag_stats, Write
|
||||
import copapy.backend as cpb
|
||||
from typing import Any
|
||||
|
||||
|
||||
def show_dag(val: value[Any]):
|
||||
out = [Store(val.net)]
|
||||
out = [Write(val.net)]
|
||||
|
||||
print(out)
|
||||
print('-- get_edges:')
|
||||
|
|
@ -26,12 +26,12 @@ def show_dag(val: value[Any]):
|
|||
print('#', p)
|
||||
|
||||
print('-- add_read_ops:')
|
||||
output_ops = list(cpb.add_load_ops(ordered_ops))
|
||||
output_ops = list(cpb.add_read_ops(ordered_ops))
|
||||
for p in output_ops:
|
||||
print('#', p)
|
||||
|
||||
print('-- add_write_ops:')
|
||||
extended_output_ops = list(cpb.add_store_ops(output_ops, const_list))
|
||||
extended_output_ops = list(cpb.add_write_ops(output_ops, const_list))
|
||||
for p in extended_output_ops:
|
||||
print('#', p)
|
||||
print('--')
|
||||
|
|
|
|||
|
|
@ -20,11 +20,7 @@ def test_fine():
|
|||
cp.cos(c_f),
|
||||
cp.tan(c_f),
|
||||
cp.abs(-c_i),
|
||||
cp.abs(-c_f),
|
||||
cp.sign(c_i),
|
||||
cp.sign(-c_f),
|
||||
cp.min(c_i, 5),
|
||||
cp.max(c_f, 5))
|
||||
cp.abs(-c_f))
|
||||
|
||||
re2_test = (a_f ** 2,
|
||||
a_i ** -1,
|
||||
|
|
@ -36,11 +32,7 @@ def test_fine():
|
|||
cp.cos(a_f),
|
||||
cp.tan(a_f),
|
||||
cp.abs(-a_i),
|
||||
cp.abs(-a_f),
|
||||
cp.sign(a_i),
|
||||
cp.sign(-a_f),
|
||||
cp.min(a_i, 5),
|
||||
cp.max(a_f, 5))
|
||||
cp.abs(-a_f))
|
||||
|
||||
ret_refe = (a_f ** 2,
|
||||
a_i ** -1,
|
||||
|
|
@ -51,12 +43,8 @@ def test_fine():
|
|||
ma.sin(a_f),
|
||||
ma.cos(a_f),
|
||||
ma.tan(a_f),
|
||||
abs(-a_i),
|
||||
abs(-a_f),
|
||||
(a_i > 0) - (a_i < 0),
|
||||
(-a_f > 0) - (-a_f < 0),
|
||||
min(a_i, 5),
|
||||
max(a_f, 5))
|
||||
cp.abs(-a_i),
|
||||
cp.abs(-a_f))
|
||||
|
||||
tg = Target()
|
||||
print('* compile and copy ...')
|
||||
|
|
@ -65,10 +53,10 @@ def test_fine():
|
|||
tg.run()
|
||||
print('* finished')
|
||||
|
||||
for test, val2, ref, name in zip(ret_test, re2_test, ret_refe, ['^2', '**-1', 'sqrt_int', 'sqrt_float', 'sin', 'cos', 'tan'] + ['other']*10):
|
||||
for test, val2, ref, name in zip(ret_test, re2_test, ret_refe, ('^2', '**-1', 'sqrt_int', 'sqrt_float', 'sin', 'cos', 'tan')):
|
||||
assert isinstance(test, cp.value)
|
||||
val = tg.read_value(test)
|
||||
print('+', name, val, ref, type(val), test.dtype)
|
||||
print('+', val, ref, type(val), test.dtype)
|
||||
#for t in (int, float, bool):
|
||||
# assert isinstance(val, t) == isinstance(ref, t), f"Result type does not match for {val} and {ref}"
|
||||
assert val == pytest.approx(ref, abs=1e-3), f"Result for {name} does not match: {val} and reference: {ref}" # pyright: ignore[reportUnknownMemberType]
|
||||
|
|
|
|||
|
|
@ -4,24 +4,26 @@ import pytest
|
|||
|
||||
def test_matrix_init():
|
||||
"""Test basic matrix initialization"""
|
||||
m1 = cp.tensor([[1, 2, 3], [4, 5, 6]])
|
||||
assert m1.shape == (2, 3)
|
||||
m1 = cp.matrix([[1, 2, 3], [4, 5, 6]])
|
||||
assert m1.rows == 2
|
||||
assert m1.cols == 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.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)
|
||||
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)
|
||||
|
||||
|
||||
def test_matrix_addition():
|
||||
"""Test matrix addition"""
|
||||
m1 = cp.tensor([[1, 2], [3, 4]])
|
||||
m2 = cp.tensor([[5, 6], [7, 8]])
|
||||
m1 = cp.matrix([[1, 2], [3, 4]])
|
||||
m2 = cp.matrix([[5, 6], [7, 8]])
|
||||
m3 = m1 + m2
|
||||
|
||||
assert m3[0] == (6, 8)
|
||||
|
|
@ -30,7 +32,7 @@ def test_matrix_addition():
|
|||
|
||||
def test_matrix_scalar_addition():
|
||||
"""Test matrix addition with scalar"""
|
||||
m1 = cp.tensor([[1, 2], [3, 4]])
|
||||
m1 = cp.matrix([[1, 2], [3, 4]])
|
||||
m2 = m1 + 5
|
||||
|
||||
assert m2[0] == (6, 7)
|
||||
|
|
@ -39,8 +41,8 @@ def test_matrix_scalar_addition():
|
|||
|
||||
def test_matrix_subtraction():
|
||||
"""Test matrix subtraction"""
|
||||
m1 = cp.tensor([[5, 6], [7, 8]])
|
||||
m2 = cp.tensor([[1, 2], [3, 4]])
|
||||
m1 = cp.matrix([[5, 6], [7, 8]])
|
||||
m2 = cp.matrix([[1, 2], [3, 4]])
|
||||
m3 = m1 - m2
|
||||
|
||||
assert m3[0] == (4, 4)
|
||||
|
|
@ -49,7 +51,7 @@ def test_matrix_subtraction():
|
|||
|
||||
def test_matrix_scalar_subtraction():
|
||||
"""Test matrix subtraction with scalar"""
|
||||
m1 = cp.tensor([[5, 6], [7, 8]])
|
||||
m1 = cp.matrix([[5, 6], [7, 8]])
|
||||
m2 = m1 - 2
|
||||
|
||||
assert m2[0] == (3, 4)
|
||||
|
|
@ -58,19 +60,17 @@ def test_matrix_scalar_subtraction():
|
|||
|
||||
def test_matrix_negation():
|
||||
"""Test matrix negation"""
|
||||
m1 = cp.tensor([[1, 2], [3, 4]])
|
||||
m1 = cp.matrix([[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.tensor([[1, 2], [3, 4]])
|
||||
m2 = cp.tensor([[5, 6], [7, 8]])
|
||||
m1 = cp.matrix([[1, 2], [3, 4]])
|
||||
m2 = cp.matrix([[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.tensor([[1, 2], [3, 4]])
|
||||
m1 = cp.matrix([[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.tensor([[6.0, 8.0], [12.0, 16.0]])
|
||||
m2 = cp.tensor([[2.0, 2.0], [3.0, 4.0]])
|
||||
m1 = cp.matrix([[6.0, 8.0], [12.0, 16.0]])
|
||||
m2 = cp.matrix([[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.tensor([[6.0, 8.0], [12.0, 16.0]])
|
||||
m1 = cp.matrix([[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,24 +109,25 @@ def test_matrix_scalar_division():
|
|||
|
||||
def test_matrix_vector_multiplication():
|
||||
"""Test matrix-vector multiplication using @ operator"""
|
||||
m = cp.tensor([[1, 2, 3], [4, 5, 6]])
|
||||
m = cp.matrix([[1, 2, 3], [4, 5, 6]])
|
||||
v = cp.vector([7, 8, 9])
|
||||
result = m @ v
|
||||
|
||||
assert isinstance(result, cp.tensor)
|
||||
assert isinstance(result, cp.vector)
|
||||
assert len(result.values) == 2
|
||||
assert result[0] == 1*7 + 2*8 + 3*9
|
||||
assert result[1] == 4*7 + 5*8 + 6*9
|
||||
assert result.values[0] == 1*7 + 2*8 + 3*9
|
||||
assert result.values[1] == 4*7 + 5*8 + 6*9
|
||||
|
||||
|
||||
def test_matrix_matrix_multiplication():
|
||||
"""Test matrix-matrix multiplication using @ operator"""
|
||||
m1 = cp.tensor([[1, 2], [3, 4]])
|
||||
m2 = cp.tensor([[5, 6], [7, 8]])
|
||||
m1 = cp.matrix([[1, 2], [3, 4]])
|
||||
m2 = cp.matrix([[5, 6], [7, 8]])
|
||||
result = m1 @ m2
|
||||
|
||||
assert isinstance(result, cp.tensor)
|
||||
assert result.shape == (2, 2)
|
||||
assert isinstance(result, cp.matrix)
|
||||
assert result.rows == 2
|
||||
assert result.cols == 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
|
||||
|
|
@ -135,10 +136,11 @@ def test_matrix_matrix_multiplication():
|
|||
|
||||
def test_matrix_transpose():
|
||||
"""Test matrix transpose"""
|
||||
m = cp.tensor([[1, 2, 3], [4, 5, 6]])
|
||||
m = cp.matrix([[1, 2, 3], [4, 5, 6]])
|
||||
mt = m.transpose()
|
||||
|
||||
assert mt.shape == (3, 2)
|
||||
assert mt.rows == 3
|
||||
assert mt.cols == 2
|
||||
assert mt[0] == (1, 4)
|
||||
assert mt[1] == (2, 5)
|
||||
assert mt[2] == (3, 6)
|
||||
|
|
@ -146,34 +148,35 @@ def test_matrix_transpose():
|
|||
|
||||
def test_matrix_transpose_property():
|
||||
"""Test matrix transpose using .T property"""
|
||||
m = cp.tensor([[1, 2, 3], [4, 5, 6]])
|
||||
m = cp.matrix([[1, 2, 3], [4, 5, 6]])
|
||||
mt = m.T
|
||||
|
||||
assert mt.shape == (3, 2)
|
||||
assert mt.rows == 3
|
||||
assert mt.cols == 2
|
||||
assert mt[0] == (1, 4)
|
||||
|
||||
|
||||
def test_matrix_row_access():
|
||||
"""Test getting a row as a vector"""
|
||||
m = cp.tensor([[1, 2, 3], [4, 5, 6]])
|
||||
row0 = m[0]
|
||||
m = cp.matrix([[1, 2, 3], [4, 5, 6]])
|
||||
row0 = m.row(0)
|
||||
|
||||
assert isinstance(row0, cp.tensor)
|
||||
assert isinstance(row0, cp.vector)
|
||||
assert row0.values == (1, 2, 3)
|
||||
|
||||
|
||||
def test_matrix_col_access():
|
||||
"""Test getting a column as a vector"""
|
||||
m = cp.tensor([[1, 2, 3], [4, 5, 6]])
|
||||
col1 = m[:, 1]
|
||||
m = cp.matrix([[1, 2, 3], [4, 5, 6]])
|
||||
col1 = m.col(1)
|
||||
|
||||
assert isinstance(col1, cp.tensor)
|
||||
assert col1 == (2, 5)
|
||||
assert isinstance(col1, cp.vector)
|
||||
assert col1.values == (2, 5)
|
||||
|
||||
|
||||
def test_matrix_trace():
|
||||
"""Test matrix trace (sum of diagonal elements)"""
|
||||
m = cp.tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
|
||||
m = cp.matrix([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
|
||||
trace = m.trace()
|
||||
|
||||
assert trace == 1 + 5 + 9
|
||||
|
|
@ -181,7 +184,7 @@ def test_matrix_trace():
|
|||
|
||||
def test_matrix_sum():
|
||||
"""Test sum of all matrix elements"""
|
||||
m = cp.tensor([[1, 2, 3], [4, 5, 6]])
|
||||
m = cp.matrix([[1, 2, 3], [4, 5, 6]])
|
||||
total = m.sum()
|
||||
|
||||
assert total == 1 + 2 + 3 + 4 + 5 + 6
|
||||
|
|
@ -189,7 +192,7 @@ def test_matrix_sum():
|
|||
|
||||
def test_matrix_map():
|
||||
"""Test mapping a function over matrix elements"""
|
||||
m = cp.tensor([[1, 2], [3, 4]])
|
||||
m = cp.matrix([[1, 2], [3, 4]])
|
||||
m_doubled = m.map(lambda x: x * 2)
|
||||
|
||||
assert m_doubled[0] == (2, 4)
|
||||
|
|
@ -198,19 +201,20 @@ def test_matrix_map():
|
|||
|
||||
def test_matrix_homogenize():
|
||||
"""Test homogenizing matrix (converting to all variables)"""
|
||||
m = cp.tensor([[1, cp.value(2)], [3, 4]])
|
||||
m = cp.matrix([[1, cp.value(2)], [3, 4]])
|
||||
m_homo = m.homogenize()
|
||||
|
||||
for row in m_homo:
|
||||
for elem in row:
|
||||
assert isinstance(elem, cp.tensor) and elem.ndim == 0
|
||||
assert isinstance(elem, cp.value)
|
||||
|
||||
|
||||
def test_identity_matrix():
|
||||
"""Test identity matrix creation"""
|
||||
m = cp.identity(3)
|
||||
|
||||
assert m.shape == (3, 3)
|
||||
assert m.rows == 3
|
||||
assert m.cols == 3
|
||||
assert m[0] == (1, 0, 0)
|
||||
assert m[1] == (0, 1, 0)
|
||||
assert m[2] == (0, 0, 1)
|
||||
|
|
@ -218,18 +222,20 @@ def test_identity_matrix():
|
|||
|
||||
def test_zeros_matrix():
|
||||
"""Test zeros matrix creation"""
|
||||
m = cp.zeros([2, 3])
|
||||
m = cp.zeros(2, 3)
|
||||
|
||||
assert m.shape == (2, 3)
|
||||
assert m.rows == 2
|
||||
assert m.cols == 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.shape == (2, 3)
|
||||
assert m.rows == 2
|
||||
assert m.cols == 3
|
||||
assert m[0] == (1, 1, 1)
|
||||
assert m[1] == (1, 1, 1)
|
||||
|
||||
|
|
@ -239,7 +245,8 @@ def test_diagonal_matrix():
|
|||
v = cp.vector([1, 2, 3])
|
||||
m = cp.diagonal(v)
|
||||
|
||||
assert m.shape == (3, 3)
|
||||
assert m.rows == 3
|
||||
assert m.cols == 3
|
||||
assert m[0] == (1, 0, 0)
|
||||
assert m[1] == (0, 2, 0)
|
||||
assert m[2] == (0, 0, 3)
|
||||
|
|
@ -247,7 +254,7 @@ def test_diagonal_matrix():
|
|||
|
||||
def test_matrix_with_variables_compiled():
|
||||
"""Test matrix operations with variables in compilation"""
|
||||
m = cp.tensor([[cp.value(1), 2], [3, cp.value(4)]])
|
||||
m = cp.matrix([[cp.value(1), 2], [3, cp.value(4)]])
|
||||
v = cp.vector([cp.value(5), 6])
|
||||
result = m @ v
|
||||
|
||||
|
|
|
|||
|
|
@ -1,5 +1,5 @@
|
|||
from copapy import NumLike, iif, value
|
||||
from copapy.backend import Store, compile_to_dag, add_read_command
|
||||
from copapy.backend import Write, compile_to_dag, add_read_command
|
||||
import subprocess
|
||||
from copapy import _binwrite
|
||||
import copapy.backend as backend
|
||||
|
|
@ -91,7 +91,7 @@ def test_compile():
|
|||
ret_test = function1(c_i) + function1(c_f) + function2(c_i) + function2(c_f) + function3(c_i) + function4(c_i) + function5(c_b) + [value(9) % 2] + iiftests(c_i) + iiftests(c_f) + [cp.asin(c_i/10)]
|
||||
ret_ref = function1(9) + function1(1.111) + function2(9) + function2(1.111) + function3(9) + function4(9) + function5(True) + [9 % 2] + iiftests(9) + iiftests(1.111) + [cp.asin(9/10)]
|
||||
|
||||
out = [Store(r) for r in ret_test]
|
||||
out = [Write(r) for r in ret_test]
|
||||
|
||||
#ret_test += [c_i, v2]
|
||||
#ret_ref += [9, 4.44, -4.44]
|
||||
|
|
|
|||
|
|
@ -1,5 +1,5 @@
|
|||
from copapy import NumLike, iif, value
|
||||
from copapy.backend import Store, compile_to_dag, add_read_command
|
||||
from copapy.backend import Write, compile_to_dag, add_read_command
|
||||
import subprocess
|
||||
from copapy import _binwrite
|
||||
import copapy.backend as backend
|
||||
|
|
@ -96,7 +96,7 @@ def test_compile():
|
|||
#ret_test = (c_i * 100 // 5, c_f * 10 // 5)
|
||||
#ret_ref = (9 * 100 // 5, 1.111 * 10 // 5)
|
||||
|
||||
out = [Store(r) for r in ret_test]
|
||||
out = [Write(r) for r in ret_test]
|
||||
|
||||
sdb = backend.stencil_db_from_package('armv6')
|
||||
dw, variables = compile_to_dag(out, sdb)
|
||||
|
|
|
|||
|
|
@ -1,5 +1,5 @@
|
|||
from copapy import NumLike, iif, value
|
||||
from copapy.backend import Store, compile_to_dag, add_read_command
|
||||
from copapy.backend import Write, compile_to_dag, add_read_command
|
||||
import subprocess
|
||||
from copapy import _binwrite
|
||||
import copapy.backend as backend
|
||||
|
|
@ -96,7 +96,7 @@ def test_compile():
|
|||
#ret_test = (c_i * 100 // 5, c_f * 10 // 5)
|
||||
#ret_ref = (9 * 100 // 5, 1.111 * 10 // 5)
|
||||
|
||||
out = [Store(r) for r in ret_test]
|
||||
out = [Write(r) for r in ret_test]
|
||||
|
||||
sdb = backend.stencil_db_from_package('armv7')
|
||||
dw, variables = compile_to_dag(out, sdb)
|
||||
|
|
|
|||
|
|
@ -1,5 +1,5 @@
|
|||
from copapy import NumLike, iif, value
|
||||
from copapy.backend import Store, compile_to_dag, add_read_command
|
||||
from copapy.backend import Write, compile_to_dag, add_read_command
|
||||
import subprocess
|
||||
from copapy import _binwrite
|
||||
import copapy.backend as backend
|
||||
|
|
@ -104,7 +104,7 @@ def test_compile():
|
|||
#ret_test = [cp.get_42(c_i)]
|
||||
#ret_ref = [cp.get_42(9)]
|
||||
|
||||
out = [Store(r) for r in ret_test]
|
||||
out = [Write(r) for r in ret_test]
|
||||
|
||||
#ret_test += [c_i, v2]
|
||||
#ret_ref += [9, 4.44, -4.44]
|
||||
|
|
@ -185,7 +185,7 @@ def test_vector_compile():
|
|||
|
||||
ret = (t2, t4, t5)
|
||||
|
||||
out = [Store(r) for r in ret]
|
||||
out = [Write(r) for r in ret]
|
||||
|
||||
sdb = backend.stencil_db_from_package('x86')
|
||||
il, variables = compile_to_dag(out, sdb)
|
||||
|
|
@ -243,7 +243,7 @@ def test_sinus():
|
|||
ret_test = [si, e]
|
||||
ret_ref = [cp.sin(a_val), (a_val + 0.87 * 2.0) ** 2 + cp.sin(a_val) + cp.sqrt(0.87)]
|
||||
|
||||
out = [Store(r) for r in ret_test]
|
||||
out = [Write(r) for r in ret_test]
|
||||
|
||||
sdb = backend.stencil_db_from_package('x86')
|
||||
dw, variables = compile_to_dag(out, sdb)
|
||||
|
|
|
|||
|
|
@ -1,257 +0,0 @@
|
|||
#!/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!")
|
||||
|
||||
144
tools/build.bat
144
tools/build.bat
|
|
@ -1,156 +1,68 @@
|
|||
@echo off
|
||||
setlocal ENABLEDELAYEDEXPANSION
|
||||
|
||||
set ARCH=%1
|
||||
if "%ARCH%"=="" set ARCH=x86_64
|
||||
|
||||
if not "%ARCH%"=="x86_64" ^
|
||||
if not "%ARCH%"=="x86" ^
|
||||
if not "%ARCH%"=="arm64" ^
|
||||
if not "%ARCH%"=="arm-v6" ^
|
||||
if not "%ARCH%"=="arm-v7" ^
|
||||
if not "%ARCH%"=="all" (
|
||||
echo Usage: %0 [x86_64^|x86^|arm64^|arm-v6^|arm-v7^|all]
|
||||
exit /b 1
|
||||
)
|
||||
|
||||
mkdir build\stencils
|
||||
mkdir build\runner
|
||||
python stencils/generate_stencils.py build/stencils/stencils.c
|
||||
|
||||
python stencils/generate_stencils.py build\stencils\stencils.c
|
||||
|
||||
REM ============================================================
|
||||
REM x86_64
|
||||
REM ============================================================
|
||||
if "%ARCH%"=="x86_64" goto BUILD_X86_64
|
||||
if "%ARCH%"=="all" goto BUILD_X86_64
|
||||
goto SKIP_X86_64
|
||||
|
||||
:BUILD_X86_64
|
||||
echo -------------x86_64 - 64 bit-----------------
|
||||
|
||||
call "C:\Program Files (x86)\Microsoft Visual Studio\2022\BuildTools\VC\Auxiliary\Build\vcvarsall.bat" x64
|
||||
|
||||
echo - Compile stencil test...
|
||||
cl /Zi /Od stencils\test.c /Fe:build\stencils\test.exe
|
||||
|
||||
echo - Build runner for Windows 64 bit...
|
||||
cl /Zi /Od /DENABLE_BASIC_LOGGING ^
|
||||
src\coparun\runmem.c ^
|
||||
src\coparun\coparun.c ^
|
||||
src\coparun\mem_man.c ^
|
||||
/Fe:build\runner\coparun.exe
|
||||
cl /Zi /Od /DENABLE_BASIC_LOGGING src\coparun\runmem.c src\coparun\coparun.c src\coparun\mem_man.c /Fe:build\runner\coparun.exe
|
||||
|
||||
echo - Build stencils for x86_64...
|
||||
REM Optimized:
|
||||
REM cl /O2 src\coparun\runmem.c src\coparun\coparun.c src\coparun\mem_man.c /Fe:build\runner\coparun.exe
|
||||
|
||||
echo - Build stencils for 64 bit...
|
||||
REM ../copapy/tools/cross_compiler_unix/packobjs.sh gcc ld ../copapy/build/musl/musl_objects_x86_64.o
|
||||
wsl gcc -fno-pic -ffunction-sections -c build/stencils/stencils.c -O3 -o build/stencils/stencils.o
|
||||
wsl ld -r build/stencils/stencils.o build/musl/musl_objects_x86_64.o -o src/copapy/obj/stencils_x86_64_O3.o
|
||||
wsl objdump -d -x src/copapy/obj/stencils_x86_64_O3.o > build/stencils/stencils_x86_64_O3.asm
|
||||
|
||||
:SKIP_X86_64
|
||||
|
||||
REM ============================================================
|
||||
REM x86 32-bit
|
||||
REM ============================================================
|
||||
if "%ARCH%"=="x86" goto BUILD_X86
|
||||
if "%ARCH%"=="all" goto BUILD_X86
|
||||
goto SKIP_X86
|
||||
|
||||
:BUILD_X86
|
||||
echo ---------------x86 - 32 bit----------------
|
||||
|
||||
echo ---------------x86 - 32 bit---------------
|
||||
call "C:\Program Files (x86)\Microsoft Visual Studio\2022\BuildTools\VC\Auxiliary\Build\vcvarsall.bat" x86
|
||||
|
||||
echo - Build runner for Windows 32 bit...
|
||||
cl /Zi /Od /DENABLE_LOGGING ^
|
||||
src\coparun\runmem.c ^
|
||||
src\coparun\coparun.c ^
|
||||
src\coparun\mem_man.c ^
|
||||
/Fe:build\runner\coparun-x86.exe
|
||||
cl /Zi /Od /DENABLE_LOGGING src\coparun\runmem.c src\coparun\coparun.c src\coparun\mem_man.c /Fe:build\runner\coparun-x86.exe
|
||||
|
||||
echo - Build runner for Linux x86 32 bit...
|
||||
wsl i686-linux-gnu-gcc-12 -static -O3 -DENABLE_LOGGING ^
|
||||
src/coparun/runmem.c ^
|
||||
src/coparun/coparun.c ^
|
||||
src/coparun/mem_man.c ^
|
||||
-o build/runner/coparun-x86
|
||||
echo - Build runner for linux x86 32 bit...
|
||||
wsl i686-linux-gnu-gcc-12 -static -Wall -Wextra -Wconversion -Wsign-conversion -Wshadow -Wstrict-overflow -O3 -DENABLE_LOGGING src/coparun/runmem.c src/coparun/coparun.c src/coparun/mem_man.c -o build/runner/coparun-x86
|
||||
|
||||
echo - Build stencils x86 32 bit...
|
||||
REM sh ../copapy/tools/cross_compiler_unix/packobjs.sh i686-linux-gnu-gcc-12 i686-linux-gnu-ld ../copapy/build/musl/musl_objects_x86.o -fno-pic
|
||||
wsl i686-linux-gnu-gcc-12 -fno-pic -ffunction-sections -c build/stencils/stencils.c -O3 -o build/stencils/stencils.o
|
||||
wsl i686-linux-gnu-ld -r build/stencils/stencils.o build/musl/musl_objects_x86.o -o src/copapy/obj/stencils_x86_O3.o
|
||||
wsl i686-linux-gnu-objdump -d -x src/copapy/obj/stencils_x86_O3.o > build/stencils/stencils_x86_O3.asm
|
||||
|
||||
:SKIP_X86
|
||||
|
||||
REM ============================================================
|
||||
REM ARM64
|
||||
REM ============================================================
|
||||
if "%ARCH%"=="arm64" goto BUILD_ARM64
|
||||
if "%ARCH%"=="all" goto BUILD_ARM64
|
||||
goto SKIP_ARM64
|
||||
|
||||
:BUILD_ARM64
|
||||
echo --------------arm64 64 bit----------------
|
||||
|
||||
wsl aarch64-linux-gnu-gcc-12 -fno-pic -ffunction-sections -c build/stencils/stencils.c -O3 -o build/stencils/stencils.o
|
||||
wsl aarch64-linux-gnu-ld -r build/stencils/stencils.o build/musl/musl_objects_arm64.o -o src/copapy/obj/stencils_arm64_O3.o
|
||||
wsl aarch64-linux-gnu-objdump -d -x src/copapy/obj/stencils_arm64_O3.o > build/stencils/stencils_arm64_O3.asm
|
||||
echo ------------------------------
|
||||
echo - Build runner for Aarch64...
|
||||
wsl aarch64-linux-gnu-gcc-12 -static -Wall -Wextra -Wconversion -Wsign-conversion -Wshadow -Wstrict-overflow -O3 -DENABLE_LOGGING src/coparun/runmem.c src/coparun/coparun.c src/coparun/mem_man.c -o build/runner/coparun-aarch64
|
||||
|
||||
echo - Build runner for AArch64...
|
||||
wsl aarch64-linux-gnu-gcc-12 -static -O3 -DENABLE_LOGGING ^
|
||||
src/coparun/runmem.c ^
|
||||
src/coparun/coparun.c ^
|
||||
src/coparun/mem_man.c ^
|
||||
-o build/runner/coparun-aarch64
|
||||
|
||||
:SKIP_ARM64
|
||||
|
||||
REM ============================================================
|
||||
REM ARM v6
|
||||
REM ============================================================
|
||||
if "%ARCH%"=="arm-v6" goto BUILD_ARMV6
|
||||
if "%ARCH%"=="all" goto BUILD_ARMV6
|
||||
goto SKIP_ARMV6
|
||||
|
||||
:BUILD_ARMV6
|
||||
echo --------------arm-v6 32 bit----------------
|
||||
|
||||
wsl arm-none-eabi-gcc -fno-pic -ffunction-sections -march=armv6 -mfpu=vfp -mfloat-abi=hard -marm ^
|
||||
-c build/stencils/stencils.c -O3 -o build/stencils/stencils.o
|
||||
|
||||
wsl arm-none-eabi-ld -r build/stencils/stencils.o build/musl/musl_objects_armv6.o ^
|
||||
$(arm-none-eabi-gcc -print-libgcc-file-name) ^
|
||||
-o src/copapy/obj/stencils_armv6_O3.o
|
||||
|
||||
REM sh ../copapy/tools/cross_compiler_unix/packobjs.sh arm-none-eabi-gcc arm-none-eabi-ld ../copapy/build/musl/musl_objects_armv6.o "-march=armv6 -mfpu=vfp -marm"
|
||||
wsl arm-none-eabi-gcc -fno-pic -ffunction-sections -march=armv6 -mfpu=vfp -mfloat-abi=hard -marm -c build/stencils/stencils.c -O3 -o build/stencils/stencils.o
|
||||
wsl arm-none-eabi-ld -r build/stencils/stencils.o build/musl/musl_objects_armv6.o $(arm-none-eabi-gcc -print-libgcc-file-name) -o src/copapy/obj/stencils_armv6_O3.o
|
||||
wsl arm-none-eabi-objdump -d -x src/copapy/obj/stencils_armv6_O3.o > build/stencils/stencils_armv6_O3.asm
|
||||
echo ------------------------------
|
||||
REM echo - Build runner
|
||||
REM wsl arm-linux-gnueabihf-gcc -march=armv6 -mfpu=vfp -marm -static -Wall -Wextra -Wconversion -Wsign-conversion -Wshadow -Wstrict-overflow -O3 -DENABLE_LOGGING src/coparun/runmem.c src/coparun/coparun.c src/coparun/mem_man.c -o build/runner/coparun-armv6
|
||||
|
||||
:SKIP_ARMV6
|
||||
|
||||
REM ============================================================
|
||||
REM ARM v7
|
||||
REM ============================================================
|
||||
if "%ARCH%"=="arm-v7" goto BUILD_ARMV7
|
||||
if "%ARCH%"=="all" goto BUILD_ARMV7
|
||||
goto END
|
||||
|
||||
:BUILD_ARMV7
|
||||
echo --------------arm-v7 32 bit----------------
|
||||
|
||||
wsl arm-none-eabi-gcc -fno-pic -ffunction-sections -march=armv7-a -mfpu=neon-vfpv3 -mfloat-abi=hard -marm ^
|
||||
-c build/stencils/stencils.c -O3 -o build/stencils/stencils.o
|
||||
|
||||
wsl arm-none-eabi-ld -r build/stencils/stencils.o build/musl/musl_objects_armv7.o ^
|
||||
$(arm-none-eabi-gcc -print-libgcc-file-name) ^
|
||||
-o src/copapy/obj/stencils_armv7_O3.o
|
||||
|
||||
REM sh ../copapy/tools/cross_compiler_unix/packobjs.sh arm-none-eabi-gcc arm-none-eabi-ld ../copapy/build/musl/musl_objects_armv7.o "-march=armv7-a -mfpu=neon-vfpv3 -marm"
|
||||
wsl arm-none-eabi-gcc -fno-pic -ffunction-sections -march=armv7-a -mfpu=neon-vfpv3 -mfloat-abi=hard -marm -c build/stencils/stencils.c -O3 -o build/stencils/stencils.o
|
||||
wsl arm-none-eabi-ld -r build/stencils/stencils.o build/musl/musl_objects_armv7.o $(arm-none-eabi-gcc -print-libgcc-file-name) -o src/copapy/obj/stencils_armv7_O3.o
|
||||
wsl arm-none-eabi-objdump -d -x src/copapy/obj/stencils_armv7_O3.o > build/stencils/stencils_armv7_O3.asm
|
||||
|
||||
echo - Build runner for ARM v7...
|
||||
wsl arm-linux-gnueabihf-gcc -static -O3 -DENABLE_LOGGING ^
|
||||
src/coparun/runmem.c ^
|
||||
src/coparun/coparun.c ^
|
||||
src/coparun/mem_man.c ^
|
||||
-o build/runner/coparun-armv7
|
||||
|
||||
:END
|
||||
echo Build completed for %ARCH%
|
||||
endlocal
|
||||
echo ------------------------------
|
||||
echo - Build runner
|
||||
wsl arm-linux-gnueabihf-gcc -march=armv7-a -mfpu=neon-vfpv3 -marm -static -Wall -Wextra -Wconversion -Wsign-conversion -Wshadow -Wstrict-overflow -O3 -DENABLE_LOGGING src/coparun/runmem.c src/coparun/coparun.c src/coparun/mem_man.c -o build/runner/coparun-armv7
|
||||
|
||||
|
|
|
|||
104
tools/build.sh
104
tools/build.sh
|
|
@ -1,16 +1,6 @@
|
|||
#!/bin/bash
|
||||
set -eux
|
||||
|
||||
ARCH=${1:-x86_64}
|
||||
|
||||
case "$ARCH" in
|
||||
(x86_64|arm-v6|arm-v7|all)
|
||||
;;
|
||||
(*)
|
||||
echo "Usage: $0 [x86_64|arm-v6|arm-v7|all]"
|
||||
exit 1
|
||||
;;
|
||||
esac
|
||||
set -e
|
||||
set -v
|
||||
|
||||
mkdir -p build/stencils
|
||||
mkdir -p build/runner
|
||||
|
|
@ -20,90 +10,34 @@ DEST=src/copapy/obj
|
|||
python3 stencils/generate_stencils.py $SRC
|
||||
mkdir -p $DEST
|
||||
|
||||
#######################################
|
||||
# x86_64
|
||||
#######################################
|
||||
if [[ "$ARCH" == "x86_64" || "$ARCH" == "all" ]]; then
|
||||
echo "--------------x86_64----------------"
|
||||
|
||||
gcc -fno-pic -ffunction-sections -c $SRC -O3 -o build/stencils/stencils.o
|
||||
ld -r build/stencils/stencils.o build/musl/musl_objects_x86_64.o \
|
||||
-o $DEST/stencils_x86_64_O3.o
|
||||
objdump -d -x $DEST/stencils_x86_64_O3.o \
|
||||
> build/stencils/stencils_x86_64_O3.asm
|
||||
ld -r build/stencils/stencils.o build/musl/musl_objects_x86_64.o -o $DEST/stencils_x86_64_O3.o
|
||||
objdump -d -x $DEST/stencils_x86_64_O3.o > build/stencils/stencils_x86_64_O3.asm
|
||||
|
||||
mkdir -p bin
|
||||
mkdir bin -p
|
||||
gcc -Wall -Wextra -Wconversion -Wsign-conversion \
|
||||
-Wshadow -Wstrict-overflow -Werror -g -O3 \
|
||||
-DENABLE_LOGGING \
|
||||
src/coparun/runmem.c \
|
||||
src/coparun/coparun.c \
|
||||
src/coparun/mem_man.c \
|
||||
-o build/runner/coparun
|
||||
fi
|
||||
src/coparun/runmem.c src/coparun/coparun.c src/coparun/mem_man.c -o build/runner/coparun
|
||||
|
||||
|
||||
#######################################
|
||||
# ARM v6
|
||||
#######################################
|
||||
if [[ "$ARCH" == "arm-v6" || "$ARCH" == "all" ]]; then
|
||||
echo "--------------arm-v6 32 bit----------------"
|
||||
|
||||
LIBGCC=$(arm-none-eabi-gcc -print-libgcc-file-name)
|
||||
#LIBM=$(arm-none-eabi-gcc -print-file-name=libm.a)
|
||||
#LIBC=$(arm-none-eabi-gcc -print-file-name=libc.a)
|
||||
|
||||
arm-none-eabi-gcc -fno-pic -ffunction-sections \
|
||||
-march=armv6 -mfpu=vfp -mfloat-abi=hard -marm \
|
||||
-c $SRC -O3 -o build/stencils/stencils.o
|
||||
arm-none-eabi-gcc -fno-pic -ffunction-sections -march=armv6 -mfpu=vfp -mfloat-abi=hard -marm -c $SRC -O3 -o build/stencils/stencils.o
|
||||
arm-none-eabi-ld -r build/stencils/stencils.o build/musl/musl_objects_armv6.o $LIBGCC -o $DEST/stencils_armv6_O3.o
|
||||
arm-none-eabi-objdump -d -x $DEST/stencils_armv6_O3.o > build/stencils/stencils_armv6_O3.asm
|
||||
arm-linux-gnueabihf-gcc -march=armv6 -mfpu=vfp -mfloat-abi=hard -marm -static -Wall -Wextra -Wconversion -Wsign-conversion -Wshadow -Wstrict-overflow -O3 -DENABLE_LOGGING src/coparun/runmem.c src/coparun/coparun.c src/coparun/mem_man.c -o build/runner/coparun-armv6
|
||||
|
||||
arm-none-eabi-ld -r \
|
||||
build/stencils/stencils.o \
|
||||
build/musl/musl_objects_armv6.o \
|
||||
$LIBGCC \
|
||||
-o $DEST/stencils_armv6_O3.o
|
||||
|
||||
arm-none-eabi-objdump -d -x \
|
||||
$DEST/stencils_armv6_O3.o \
|
||||
> build/stencils/stencils_armv6_O3.asm
|
||||
|
||||
arm-linux-gnueabihf-gcc \
|
||||
-march=armv6 -mfpu=vfp -mfloat-abi=hard -marm -static \
|
||||
-Wall -Wextra -Wconversion -Wsign-conversion \
|
||||
-Wshadow -Wstrict-overflow -O3 \
|
||||
-DENABLE_LOGGING \
|
||||
src/coparun/runmem.c \
|
||||
src/coparun/coparun.c \
|
||||
src/coparun/mem_man.c \
|
||||
-o build/runner/coparun-armv6
|
||||
fi
|
||||
|
||||
#######################################
|
||||
# ARM v7
|
||||
#######################################
|
||||
if [[ "$ARCH" == "arm-v7" || "$ARCH" == "all" ]]; then
|
||||
echo "--------------arm-v7 32 bit----------------"
|
||||
|
||||
LIBGCC=$(arm-none-eabi-gcc -print-libgcc-file-name)
|
||||
#LIBM=$(arm-none-eabi-gcc -print-file-name=libm.a)
|
||||
#LIBC=$(arm-none-eabi-gcc -print-file-name=libc.a)
|
||||
|
||||
arm-none-eabi-gcc -fno-pic -ffunction-sections \
|
||||
-march=armv7-a -mfpu=neon-vfpv3 -mfloat-abi=hard -marm \
|
||||
-c $SRC -O3 -o build/stencils/stencils.o
|
||||
|
||||
arm-none-eabi-ld -r \
|
||||
build/stencils/stencils.o \
|
||||
build/musl/musl_objects_armv7.o \
|
||||
$LIBGCC \
|
||||
-o $DEST/stencils_armv7_O3.o
|
||||
|
||||
arm-none-eabi-objdump -d -x \
|
||||
$DEST/stencils_armv7_O3.o \
|
||||
> build/stencils/stencils_armv7_O3.asm
|
||||
|
||||
arm-linux-gnueabihf-gcc \
|
||||
-march=armv7-a -mfpu=neon-vfpv3 -mfloat-abi=hard -marm -static \
|
||||
-Wall -Wextra -Wconversion -Wsign-conversion \
|
||||
-Wshadow -Wstrict-overflow -O3 \
|
||||
-DENABLE_LOGGING \
|
||||
src/coparun/runmem.c \
|
||||
src/coparun/coparun.c \
|
||||
src/coparun/mem_man.c \
|
||||
-o build/runner/coparun-armv7
|
||||
fi
|
||||
arm-none-eabi-gcc -fno-pic -ffunction-sections -march=armv7-a -mfpu=neon-vfpv3 -mfloat-abi=hard -marm -c $SRC -O3 -o build/stencils/stencils.o
|
||||
arm-none-eabi-ld -r build/stencils/stencils.o build/musl/musl_objects_armv7.o $LIBGCC -o $DEST/stencils_armv7_O3.o
|
||||
arm-none-eabi-objdump -d -x $DEST/stencils_armv7_O3.o > build/stencils/stencils_armv7_O3.asm
|
||||
arm-linux-gnueabihf-gcc -march=armv7-a -mfpu=neon-vfpv3 -mfloat-abi=hard -marm -static -Wall -Wextra -Wconversion -Wsign-conversion -Wshadow -Wstrict-overflow -O3 -DENABLE_LOGGING src/coparun/runmem.c src/coparun/coparun.c src/coparun/mem_man.c -o build/runner/coparun-armv7
|
||||
|
|
|
|||
|
|
@ -11,30 +11,20 @@ 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,11 +41,6 @@ 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
|
||||
|
||||
|
|
|
|||
|
|
@ -10,6 +10,18 @@ if [[ "${GITHUB_REF:-}" == refs/tags/* ]]; then
|
|||
exit 0
|
||||
fi
|
||||
|
||||
# Detected version from Git
|
||||
TAG_NAME=$(git describe --tags --abbrev=0)
|
||||
echo "version=$TAG_NAME" >> "${GITHUB_OUTPUT:-/dev/stdout}"
|
||||
# Otherwise, extract version from pyproject.toml
|
||||
if [[ -f pyproject.toml ]]; then
|
||||
VERSION=$(grep -E '^version\s*=' pyproject.toml \
|
||||
| sed -E 's/version\s*=\s*"([^"]+)"/\1/' \
|
||||
| tr -d '\r\n')
|
||||
if [[ -z "$VERSION" ]]; then
|
||||
echo "! Could not find version in pyproject.toml" >&2
|
||||
exit 1
|
||||
fi
|
||||
echo "Detected version from pyproject.toml: v$VERSION-beta"
|
||||
echo "version=v$VERSION-beta" >> "${GITHUB_OUTPUT:-/dev/stdout}"
|
||||
else
|
||||
echo "! pyproject.toml not found" >&2
|
||||
exit 1
|
||||
fi
|
||||
|
|
|
|||
|
|
@ -1,5 +1,5 @@
|
|||
from copapy import value
|
||||
from copapy.backend import Store, compile_to_dag, stencil_db_from_package
|
||||
from copapy.backend import Write, compile_to_dag, stencil_db_from_package
|
||||
from copapy._binwrite import Command
|
||||
|
||||
input = value(9.0)
|
||||
|
|
@ -8,7 +8,7 @@ result = input ** 2 / 3.3 + 5
|
|||
|
||||
arch = 'native'
|
||||
sdb = stencil_db_from_package(arch)
|
||||
dw, _ = compile_to_dag([Store(result)], sdb)
|
||||
dw, _ = compile_to_dag([Write(result)], sdb)
|
||||
|
||||
# Instruct runner to dump patched code to a file:
|
||||
dw.write_com(Command.DUMP_CODE)
|
||||
|
|
|
|||
Loading…
Reference in New Issue