rsnumpy 1.1.2__tar.gz → 1.1.3__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (48) hide show
  1. {rsnumpy-1.1.2 → rsnumpy-1.1.3}/.gitignore +1 -0
  2. {rsnumpy-1.1.2 → rsnumpy-1.1.3}/Cargo.lock +54 -1
  3. rsnumpy-1.1.3/Cargo.toml +49 -0
  4. {rsnumpy-1.1.2 → rsnumpy-1.1.3}/PKG-INFO +1 -1
  5. {rsnumpy-1.1.2 → rsnumpy-1.1.3}/build_wheel.sh +8 -8
  6. {rsnumpy-1.1.2 → rsnumpy-1.1.3}/pyproject.toml +1 -1
  7. {rsnumpy-1.1.2 → rsnumpy-1.1.3}/python/rsnumpy/__init__.py +13 -2
  8. {rsnumpy-1.1.2 → rsnumpy-1.1.3}/python/rsnumpy/array_methods.py +2 -1
  9. {rsnumpy-1.1.2 → rsnumpy-1.1.3}/python/rsnumpy/linalg/__init__.py +14 -3
  10. rsnumpy-1.1.3/src/arithmetic/mod.rs +47 -0
  11. rsnumpy-1.1.3/src/bitwise/mod.rs +92 -0
  12. rsnumpy-1.1.3/src/buffer/mod.rs +271 -0
  13. rsnumpy-1.1.3/src/creation/mod.rs +284 -0
  14. rsnumpy-1.1.3/src/formatting/mod.rs +229 -0
  15. {rsnumpy-1.1.2 → rsnumpy-1.1.3}/src/indexing.rs +88 -47
  16. rsnumpy-1.1.3/src/io/mod.rs +395 -0
  17. rsnumpy-1.1.3/src/lib.rs +2025 -0
  18. {rsnumpy-1.1.2 → rsnumpy-1.1.3}/src/linalg.rs +58 -186
  19. rsnumpy-1.1.3/src/logic/mod.rs +167 -0
  20. rsnumpy-1.1.3/src/manipulation/mod.rs +735 -0
  21. rsnumpy-1.1.3/src/mathematics/mod.rs +309 -0
  22. rsnumpy-1.1.3/src/poly/mod.rs +312 -0
  23. rsnumpy-1.1.3/src/searching/mod.rs +304 -0
  24. rsnumpy-1.1.3/src/setops/mod.rs +433 -0
  25. rsnumpy-1.1.3/src/sorting/mod.rs +346 -0
  26. rsnumpy-1.1.3/src/statistics/mod.rs +883 -0
  27. {rsnumpy-1.1.2 → rsnumpy-1.1.3}/test/run_test.py +105 -0
  28. rsnumpy-1.1.3/test/simple.csv +5 -0
  29. rsnumpy-1.1.3/test/test_new.txt +3777 -0
  30. rsnumpy-1.1.2/Cargo.toml +0 -42
  31. rsnumpy-1.1.2/src/lib.rs +0 -6736
  32. {rsnumpy-1.1.2 → rsnumpy-1.1.3}/.github/workflows/ci.yml +0 -0
  33. {rsnumpy-1.1.2 → rsnumpy-1.1.3}/.github/workflows/release.yml +0 -0
  34. {rsnumpy-1.1.2 → rsnumpy-1.1.3}/LICENSE +0 -0
  35. {rsnumpy-1.1.2 → rsnumpy-1.1.3}/README.md +0 -0
  36. {rsnumpy-1.1.2 → rsnumpy-1.1.3}/RELEASE_NOTES.md +0 -0
  37. {rsnumpy-1.1.2 → rsnumpy-1.1.3}/python/rsnumpy/_extra.py +0 -0
  38. {rsnumpy-1.1.2 → rsnumpy-1.1.3}/python/rsnumpy/array_ops.py +0 -0
  39. {rsnumpy-1.1.2 → rsnumpy-1.1.3}/python/rsnumpy/char.py +0 -0
  40. {rsnumpy-1.1.2 → rsnumpy-1.1.3}/python/rsnumpy/io.py +0 -0
  41. {rsnumpy-1.1.2 → rsnumpy-1.1.3}/python/rsnumpy/math_functions.py +0 -0
  42. {rsnumpy-1.1.2 → rsnumpy-1.1.3}/python/rsnumpy/matlib.py +0 -0
  43. {rsnumpy-1.1.2 → rsnumpy-1.1.3}/python/rsnumpy/polynomial/__init__.py +0 -0
  44. {rsnumpy-1.1.2 → rsnumpy-1.1.3}/python/rsnumpy/random/__init__.py +0 -0
  45. {rsnumpy-1.1.2 → rsnumpy-1.1.3}/python/rsnumpy/statistics.py +0 -0
  46. {rsnumpy-1.1.2 → rsnumpy-1.1.3}/rust-toolchain.toml +0 -0
  47. {rsnumpy-1.1.2 → rsnumpy-1.1.3}/src/fft.rs +0 -0
  48. {rsnumpy-1.1.2 → rsnumpy-1.1.3}/src/random.rs +0 -0
@@ -1,5 +1,6 @@
1
1
  # Rust
2
2
  target/
3
+ _split_tool/
3
4
  wheelhouse/
4
5
  Cargo.lock
5
6
  rsnum*.whl
@@ -2,6 +2,12 @@
2
2
  # It is not intended for manual editing.
3
3
  version = 4
4
4
 
5
+ [[package]]
6
+ name = "accelerate-src"
7
+ version = "0.3.2"
8
+ source = "registry+https://github.com/rust-lang/crates.io-index"
9
+ checksum = "415ed64958754dbe991900f3940677e6a7eefb4d7367afd70d642677b0c7d19d"
10
+
5
11
  [[package]]
6
12
  name = "adler2"
7
13
  version = "2.0.1"
@@ -25,6 +31,15 @@ version = "1.5.1"
25
31
  source = "registry+https://github.com/rust-lang/crates.io-index"
26
32
  checksum = "f2032f911046de80f0a198e0901378627c33f59ea0ac00e363d481118bd70a53"
27
33
 
34
+ [[package]]
35
+ name = "blas-src"
36
+ version = "0.10.0"
37
+ source = "registry+https://github.com/rust-lang/crates.io-index"
38
+ checksum = "b95e83dc868db96e69795c0213143095f03de9dd3252f205d4ac716e4076a7e0"
39
+ dependencies = [
40
+ "accelerate-src",
41
+ ]
42
+
28
43
  [[package]]
29
44
  name = "block-buffer"
30
45
  version = "0.12.1"
@@ -140,6 +155,15 @@ dependencies = [
140
155
  "cfg-if",
141
156
  ]
142
157
 
158
+ [[package]]
159
+ name = "crossbeam-channel"
160
+ version = "0.5.16"
161
+ source = "registry+https://github.com/rust-lang/crates.io-index"
162
+ checksum = "d85363c37faeca707aef026efa9f3b34d077bce547e48f770770625c6013679e"
163
+ dependencies = [
164
+ "crossbeam-utils",
165
+ ]
166
+
143
167
  [[package]]
144
168
  name = "crossbeam-deque"
145
169
  version = "0.8.7"
@@ -286,6 +310,12 @@ version = "0.5.0"
286
310
  source = "registry+https://github.com/rust-lang/crates.io-index"
287
311
  checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea"
288
312
 
313
+ [[package]]
314
+ name = "hermit-abi"
315
+ version = "0.5.2"
316
+ source = "registry+https://github.com/rust-lang/crates.io-index"
317
+ checksum = "fc0fef456e4baa96da950455cd02c081ca953b141298e41db3fc7e36b1da849c"
318
+
289
319
  [[package]]
290
320
  name = "hmac"
291
321
  version = "0.13.0"
@@ -384,7 +414,10 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
384
414
  checksum = "a06de3016e9fae57a36fd14dba131fccf49f74b40b7fbdb472f96e361ec71a08"
385
415
  dependencies = [
386
416
  "autocfg",
417
+ "num_cpus",
418
+ "once_cell",
387
419
  "rawpointer",
420
+ "thread-tree",
388
421
  ]
389
422
 
390
423
  [[package]]
@@ -456,6 +489,16 @@ dependencies = [
456
489
  "libm",
457
490
  ]
458
491
 
492
+ [[package]]
493
+ name = "num_cpus"
494
+ version = "1.17.0"
495
+ source = "registry+https://github.com/rust-lang/crates.io-index"
496
+ checksum = "91df4bbde75afed763b708b7eee1e8e7651e02d97f6d5dd763e89367e957b23b"
497
+ dependencies = [
498
+ "hermit-abi",
499
+ "libc",
500
+ ]
501
+
459
502
  [[package]]
460
503
  name = "once_cell"
461
504
  version = "1.21.4"
@@ -656,8 +699,9 @@ dependencies = [
656
699
 
657
700
  [[package]]
658
701
  name = "rsnumpy"
659
- version = "1.1.2"
702
+ version = "1.1.3"
660
703
  dependencies = [
704
+ "blas-src",
661
705
  "ndarray",
662
706
  "num-complex",
663
707
  "num-traits",
@@ -802,6 +846,15 @@ dependencies = [
802
846
  "syn",
803
847
  ]
804
848
 
849
+ [[package]]
850
+ name = "thread-tree"
851
+ version = "0.3.3"
852
+ source = "registry+https://github.com/rust-lang/crates.io-index"
853
+ checksum = "ffbd370cb847953a25954d9f63e14824a36113f8c72eecf6eccef5dc4b45d630"
854
+ dependencies = [
855
+ "crossbeam-channel",
856
+ ]
857
+
805
858
  [[package]]
806
859
  name = "time"
807
860
  version = "0.3.53"
@@ -0,0 +1,49 @@
1
+ [package]
2
+ name = "rsnumpy"
3
+ version = "1.1.3"
4
+ edition = "2024"
5
+ readme = "README.md"
6
+
7
+ [lib]
8
+ name = "rsnumpy"
9
+ crate-type = ["cdylib"]
10
+
11
+ [dependencies]
12
+ pyo3 = { version = "0.29.0", features = ["extension-module"] }
13
+ # matrixmultiply-threading:非 BLAS 平台上让纯 Rust GEMM 内核多线程并行。
14
+ ndarray = { version = "0.17.2", features = [
15
+ "rayon",
16
+ "serde",
17
+ "matrixmultiply-threading",
18
+ ] }
19
+ rayon = "1.12.0"
20
+ rand = "0.10.2"
21
+ rand_distr = "0.6.0"
22
+ num-traits = "0.2.19"
23
+ rustfft = "6.4.1"
24
+ zip = "8.6.0"
25
+ num-complex = "0.4.6"
26
+ thiserror = "2.0.0"
27
+
28
+ # macOS:链接系统自带的 Accelerate 框架作为 BLAS 后端(无需额外安装),
29
+ # matmul/dot 由 ndarray 的 blas 特性自动分派到原生高性能 GEMM。
30
+ # 其它平台不启用 BLAS,自动回退到上面的 matrixmultiply 多线程纯 Rust 内核。
31
+ [target.'cfg(target_os = "macos")'.dependencies]
32
+ ndarray = { version = "0.17.2", features = ["blas"] }
33
+ blas-src = { version = "0.10", features = ["accelerate"] }
34
+
35
+ [profile.release]
36
+ opt-level = 3
37
+ lto = "fat"
38
+ codegen-units = 1
39
+ strip = "symbols"
40
+
41
+ [profile.dev]
42
+ opt-level = 0
43
+ debug = true
44
+
45
+ [lints.clippy]
46
+ too_many_arguments = "allow"
47
+ type_complexity = "allow"
48
+ cast_precision_loss = "allow"
49
+ cast_sign_loss = "allow"
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: rsnumpy
3
- Version: 1.1.2
3
+ Version: 1.1.3
4
4
  Classifier: Programming Language :: Rust
5
5
  Classifier: Programming Language :: Python :: 3
6
6
  Classifier: Programming Language :: Python :: 3 :: Only
@@ -24,7 +24,7 @@ rm -f python/rsnumpy/__init__.py.bak
24
24
  echo " -> python/rsnumpy/__init__.py __version__ updated to $VERSION"
25
25
 
26
26
  # ========== 检查 cargo 是否可用(maturin 构建需要) ==========
27
- if ! command -v cargo >/dev/null 2>&1; then
27
+ if ! command -v cargo; then
28
28
  echo "Error: 'cargo' not found in PATH." >&2
29
29
  echo "Please install Rust and Cargo from https://rustup.rs/" >&2
30
30
  echo "After installation, run: source \$HOME/.cargo/env" >&2
@@ -38,10 +38,10 @@ if [[ -f rust-toolchain.toml ]]; then
38
38
  RUST_CHANNEL=$(grep -m1 '^channel *= *' rust-toolchain.toml | sed -E 's/.*"([^"]+)".*/\1/')
39
39
  if [[ -n "${RUST_CHANNEL:-}" ]]; then
40
40
  echo "Pinned Rust toolchain (rust-toolchain.toml): $RUST_CHANNEL"
41
- if command -v rustup >/dev/null 2>&1; then
41
+ if command -v rustup; then
42
42
  # 幂等:已安装则为空操作;离线且缺失时失败被忽略,交由后续 clippy 步骤明确报错。
43
- rustup toolchain install "$RUST_CHANNEL" >/dev/null 2>&1 || true
44
- rustup component add clippy --toolchain "$RUST_CHANNEL" >/dev/null 2>&1 || true
43
+ rustup toolchain install "$RUST_CHANNEL" || true
44
+ rustup component add clippy --toolchain "$RUST_CHANNEL" || true
45
45
  else
46
46
  echo " -> rustup not found; using the cargo/clippy already on PATH." >&2
47
47
  fi
@@ -88,15 +88,15 @@ if [[ -f "$VENV_PY" && -z "$PYTHON_EXEC_SET" ]]; then
88
88
  fi
89
89
 
90
90
  # Ensure chosen Python executable exists or is runnable
91
- if ! command -v "$PYTHON_EXEC" >/dev/null 2>&1 && [[ ! -x "$PYTHON_EXEC" && ! -f "$PYTHON_EXEC" ]]; then
91
+ if ! command -v "$PYTHON_EXEC" && [[ ! -x "$PYTHON_EXEC" && ! -f "$PYTHON_EXEC" ]]; then
92
92
  echo "Error: Python executable '$PYTHON_EXEC' not found or not executable." >&2
93
93
  exit 1
94
94
  fi
95
95
 
96
96
  # Prefer a maturin on PATH, otherwise try running maturin via the chosen Python
97
- if command -v maturin >/dev/null 2>&1; then
97
+ if command -v maturin; then
98
98
  MATURIN_MODE="path"
99
- elif "$PYTHON_EXEC" -c "import maturin" >/dev/null 2>&1; then
99
+ elif "$PYTHON_EXEC" -c "import maturin"; then
100
100
  MATURIN_MODE="python-module"
101
101
  else
102
102
  echo "Error: maturin not found. Install it in your chosen Python (e.g. '$PYTHON_EXEC -m pip install maturin')." >&2
@@ -139,7 +139,7 @@ if [[ -n "$WHEEL" && -f "$WHEEL" ]]; then
139
139
  echo "Installing into .venv ..."
140
140
  # --no-deps:本地 wheel 已含全部绑定代码,无需重新解析依赖。
141
141
  # 优先用 uv(显式 --python 指向本 venv,避免依赖是否已 activate);无 uv 时回退到 pip。
142
- if command -v uv >/dev/null 2>&1; then
142
+ if command -v uv; then
143
143
  uv pip install --python "$VENV_PY" --reinstall --no-deps "$WHEEL"
144
144
  else
145
145
  "$VENV_PY" -m pip install --force-reinstall --no-deps "$WHEEL"
@@ -4,7 +4,7 @@ build-backend = "maturin"
4
4
 
5
5
  [project]
6
6
  name = "rsnumpy"
7
- version = "1.1.2"
7
+ version = "1.1.3"
8
8
  requires-python = ">=3.10"
9
9
  description = "A Rust-powered NumPy-compatible array library."
10
10
  readme = "README.md"
@@ -32,7 +32,7 @@ from .random import random_module as _random_module
32
32
  from . import char as _char_module
33
33
  from . import matlib as _matlib_module
34
34
 
35
- __version__ = "1.1.2"
35
+ __version__ = "1.1.3"
36
36
 
37
37
  # 捕获内建函数别名:_extra 挂载会向本模块 globals 注入同名的 numpy 函数
38
38
  # (all/any/round),会遮蔽内建函数。以下别名保证本文件内部逻辑始终使用内建实现。
@@ -2273,7 +2273,18 @@ def ix_(*args):
2273
2273
  >>> x = np.arange(32).reshape((8, 4))
2274
2274
  >>> x[np.ix_([1,5,7,2], [0,3,1,2])]
2275
2275
  """
2276
- raw = _core.ix_rs(args)
2276
+ # Rust 层的 ix_rs 只解析 Python list/tuple,故先把 ndarray 等参数
2277
+ # 归一化为一维下标列表;布尔序列按 numpy 语义取 True 位置的下标。
2278
+ prepared = []
2279
+ for a in args:
2280
+ dtype = getattr(a, '_dtype', None)
2281
+ seq = a.tolist() if hasattr(a, 'tolist') else list(a)
2282
+ if dtype == 'bool' or (len(seq) > 0 and _py_all(isinstance(v, bool) for v in seq)):
2283
+ seq = [i for i, v in enumerate(seq) if v]
2284
+ else:
2285
+ seq = [int(v) for v in seq]
2286
+ prepared.append(seq)
2287
+ raw = _core.ix_rs(tuple(prepared))
2277
2288
  return tuple(ndarray._wrap(r) for r in raw)
2278
2289
 
2279
2290
 
@@ -21,7 +21,8 @@ class NdArrayMethods:
21
21
  @staticmethod
22
22
  def astype(arr, dtype):
23
23
  """将数组转换为指定的数据类型。"""
24
- return _wrap_result(arr._array.astype(dtype))
24
+ name = dtype if isinstance(dtype, str) else getattr(dtype, "__name__", str(dtype))
25
+ return _wrap_result(arr._array.astype(name), dtype=name)
25
26
 
26
27
  @staticmethod
27
28
  def reshape(arr, *shape):
@@ -52,7 +52,6 @@ class linalg_module:
52
52
  """计算两个数组的点积。"""
53
53
  from ..__init__ import ndarray
54
54
  result = _core.linalg.dot(_ensure(a), _ensure(b))
55
- result = _round_array(result)
56
55
  a_dtype = getattr(a, '_dtype', 'float64')
57
56
  b_dtype = getattr(b, '_dtype', 'float64')
58
57
  if a_dtype in ('float64', 'float32', 'float16') or b_dtype in ('float64', 'float32', 'float16'):
@@ -68,12 +67,24 @@ class linalg_module:
68
67
  @staticmethod
69
68
  def inner(a, b):
70
69
  """计算两个数组的内积。"""
71
- return _wrap(_core.linalg.inner(_ensure(a), _ensure(b)))
70
+ from ..__init__ import ndarray
71
+ result = _core.linalg.inner(_ensure(a), _ensure(b))
72
+ a_dtype = getattr(a, '_dtype', 'float64')
73
+ b_dtype = getattr(b, '_dtype', 'float64')
74
+ if a_dtype in ('float64', 'float32', 'float16') or b_dtype in ('float64', 'float32', 'float16'):
75
+ return ndarray._wrap(result, _dtype='float64')
76
+ return ndarray._wrap(result, _dtype='int64')
72
77
 
73
78
  @staticmethod
74
79
  def matmul(a, b):
75
80
  """计算两个数组的矩阵乘积。"""
76
- return _wrap(_core.linalg.matmul(_ensure(a), _ensure(b)))
81
+ from ..__init__ import ndarray
82
+ result = _core.linalg.matmul(_ensure(a), _ensure(b))
83
+ a_dtype = getattr(a, '_dtype', 'float64')
84
+ b_dtype = getattr(b, '_dtype', 'float64')
85
+ if a_dtype in ('float64', 'float32', 'float16') or b_dtype in ('float64', 'float32', 'float16'):
86
+ return ndarray._wrap(result, _dtype='float64')
87
+ return ndarray._wrap(result, _dtype='int64')
77
88
 
78
89
  @staticmethod
79
90
  def inv(a):
@@ -0,0 +1,47 @@
1
+ use crate::*;
2
+
3
+ #[pyfunction]
4
+ fn add(x1: &NdArray, x2: &NdArray) -> PyResult<NdArray> {
5
+ let result = broadcast_binary_op(&x1.data, &x2.data, |a, b| a + b)?;
6
+ Ok(NdArray { data: result })
7
+ }
8
+
9
+ #[pyfunction]
10
+ fn subtract(x1: &NdArray, x2: &NdArray) -> PyResult<NdArray> {
11
+ let result = broadcast_binary_op(&x1.data, &x2.data, |a, b| a - b)?;
12
+ Ok(NdArray { data: result })
13
+ }
14
+
15
+ #[pyfunction]
16
+ fn multiply(x1: &NdArray, x2: &NdArray) -> PyResult<NdArray> {
17
+ let result = broadcast_binary_op(&x1.data, &x2.data, |a, b| a * b)?;
18
+ Ok(NdArray { data: result })
19
+ }
20
+
21
+ #[pyfunction]
22
+ fn divide(x1: &NdArray, x2: &NdArray) -> PyResult<NdArray> {
23
+ let result = broadcast_binary_op(&x1.data, &x2.data, |a, b| a / b)?;
24
+ Ok(NdArray { data: result })
25
+ }
26
+
27
+ #[pyfunction]
28
+ fn power(x1: &NdArray, x2: &NdArray) -> PyResult<NdArray> {
29
+ let result = broadcast_binary_op(&x1.data, &x2.data, |a, b| a.powf(b))?;
30
+ Ok(NdArray { data: result })
31
+ }
32
+
33
+ #[pyfunction]
34
+ fn mod_(x1: &NdArray, x2: &NdArray) -> PyResult<NdArray> {
35
+ let result = broadcast_binary_op(&x1.data, &x2.data, |a, b| a % b)?;
36
+ Ok(NdArray { data: result })
37
+ }
38
+
39
+ pub(crate) fn register(m: &Bound<'_, PyModule>) -> PyResult<()> {
40
+ m.add_function(wrap_pyfunction!(add, m)?)?;
41
+ m.add_function(wrap_pyfunction!(subtract, m)?)?;
42
+ m.add_function(wrap_pyfunction!(multiply, m)?)?;
43
+ m.add_function(wrap_pyfunction!(divide, m)?)?;
44
+ m.add_function(wrap_pyfunction!(power, m)?)?;
45
+ m.add_function(wrap_pyfunction!(mod_, m)?)?;
46
+ Ok(())
47
+ }
@@ -0,0 +1,92 @@
1
+ use crate::*;
2
+
3
+ #[pyfunction]
4
+ fn invert(x: &NdArray) -> NdArray {
5
+ NdArray {
6
+ data: x.data.mapv(|v| if v == 0.0 { 1.0 } else { 0.0 }),
7
+ }
8
+ }
9
+
10
+ #[pyfunction]
11
+ fn bitwise_and(x1: &NdArray, x2: &NdArray) -> PyResult<NdArray> {
12
+ let result = broadcast_binary_op(&x1.data, &x2.data, |a, b| {
13
+ let bits = (a as i64) & (b as i64);
14
+ bits as f64
15
+ })?;
16
+ Ok(NdArray { data: result })
17
+ }
18
+
19
+ #[pyfunction]
20
+ fn bitwise_or(x1: &NdArray, x2: &NdArray) -> PyResult<NdArray> {
21
+ let result = broadcast_binary_op(&x1.data, &x2.data, |a, b| {
22
+ let bits = (a as i64) | (b as i64);
23
+ bits as f64
24
+ })?;
25
+ Ok(NdArray { data: result })
26
+ }
27
+
28
+ #[pyfunction]
29
+ fn bitwise_xor(x1: &NdArray, x2: &NdArray) -> PyResult<NdArray> {
30
+ let result = broadcast_binary_op(&x1.data, &x2.data, |a, b| {
31
+ let bits = (a as i64) ^ (b as i64);
32
+ bits as f64
33
+ })?;
34
+ Ok(NdArray { data: result })
35
+ }
36
+
37
+ #[pyfunction]
38
+ fn left_shift(x1: &NdArray, x2: &NdArray) -> PyResult<NdArray> {
39
+ let result = broadcast_binary_op(&x1.data, &x2.data, |a, b| {
40
+ let bits = (a as i64) << (b as i64);
41
+ bits as f64
42
+ })?;
43
+ Ok(NdArray { data: result })
44
+ }
45
+
46
+ #[pyfunction]
47
+ fn right_shift(x1: &NdArray, x2: &NdArray) -> PyResult<NdArray> {
48
+ let result = broadcast_binary_op(&x1.data, &x2.data, |a, b| {
49
+ let bits = (a as i64) >> (b as i64);
50
+ bits as f64
51
+ })?;
52
+ Ok(NdArray { data: result })
53
+ }
54
+
55
+ #[pyfunction]
56
+ fn bitwise_not(x: &NdArray) -> NdArray {
57
+ NdArray {
58
+ data: x.data.mapv(|v| {
59
+ let bits = !(v as i64);
60
+ bits as f64
61
+ }),
62
+ }
63
+ }
64
+
65
+ #[pyfunction]
66
+ fn binary_repr(num: i64, width: Option<usize>) -> String {
67
+ let w = if let Some(w) = width {
68
+ w
69
+ } else {
70
+ let bits = 64 - num.abs().leading_zeros();
71
+ if bits == 0 { 1 } else { bits as usize }
72
+ };
73
+
74
+ if num >= 0 {
75
+ format!("{num:0w$b}")
76
+ } else {
77
+ let val = (1 << w) + num;
78
+ format!("{val:0w$b}")
79
+ }
80
+ }
81
+
82
+ pub(crate) fn register(m: &Bound<'_, PyModule>) -> PyResult<()> {
83
+ m.add_function(wrap_pyfunction!(invert, m)?)?;
84
+ m.add_function(wrap_pyfunction!(bitwise_and, m)?)?;
85
+ m.add_function(wrap_pyfunction!(bitwise_or, m)?)?;
86
+ m.add_function(wrap_pyfunction!(bitwise_xor, m)?)?;
87
+ m.add_function(wrap_pyfunction!(bitwise_not, m)?)?;
88
+ m.add_function(wrap_pyfunction!(left_shift, m)?)?;
89
+ m.add_function(wrap_pyfunction!(right_shift, m)?)?;
90
+ m.add_function(wrap_pyfunction!(binary_repr, m)?)?;
91
+ Ok(())
92
+ }