tensorstudio 1.3.3__tar.gz → 1.3.5__tar.gz

{tensorstudio-1.3.3 → tensorstudio-1.3.5}/CHANGELOG.md

@@ -2,6 +2,28 @@
 
 ## Unreleased
 
+## 1.3.5 - 2026-07-07
+
+- Completed the next ordered correctness-roadmap item with native C++
+  `argmax` and `argmin` operations.
+- Added Tensor methods, functional ops, top-level Python exports, and type
+  stubs for arg reductions.
+- Added NumPy parity tests for all-element and axis arg reductions, including
+  `keepdims`, negative axes, first-tie behavior, dtype, and error cases.
+- Documented arg-reduction semantics and limitations.
+
+## 1.3.4 - 2026-07-07
+
+- Completed the next ordered correctness-roadmap item with tuple/list-axis
+  reductions for `sum`, `mean`, `max`, and `min`.
+- Added C++ binding-layer axis normalization for reduction methods and
+  functional ops, including negative-axis handling, duplicate-axis errors, and
+  `axis=()` no-op behavior.
+- Added NumPy parity tests and autograd coverage for tuple-axis reductions.
+- Added ordered roadmap entries for supported ONNX/model metadata extraction,
+  supported ONNX import/execution, safe model-format research, and image/tensor
+  dataset creation helpers.
+
 ## 1.3.3 - 2026-07-07
 
 - Hardened CI type checks for newer Python/mypy combinations by annotating

{tensorstudio-1.3.3 → tensorstudio-1.3.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: tensorstudio
-Version: 1.3.3
+Version: 1.3.5
 Summary: TensorStudio is a compact C++ tensor and autograd engine with a Python API for learning, experimentation, and lightweight ML workloads.
 Keywords: tensor,autograd,machine-learning,cpp,pybind11
 Author: TensorStudio contributors
@@ -71,7 +71,7 @@ Description-Content-Type: text/markdown
 TensorStudio is a compact C++ tensor and autograd engine with a Python API for
 learning, experimentation, and lightweight ML workloads.
 
-TensorStudio `1.3.3` is a CPU-only stable API foundation. It is eager-only,
+TensorStudio `1.3.5` is a CPU-only stable API foundation. It is eager-only,
 intentionally small, and not a replacement for mature ML frameworks.
 
 ## Install
@@ -472,8 +472,10 @@ tokens or print secrets.
   pretrained model zoos, detection/segmentation training stacks, video IO, or
   GPU image kernels yet.
 - ONNX support is export-only for a limited set of TensorStudio modules.
-- Reductions support all-element or single-axis reductions, not tuple-axis
-  reductions yet.
+- Reductions support all-element, single-axis, and tuple/list-axis reductions
+  for `sum`, `mean`, `max`, and `min`.
+- Arg reductions support all-element flat indices or one axis at a time for
+  `argmax` and `argmin`.
 - No sparse tensors or advanced indexing.
 - Dtype casting is basic and does not include a full promotion/casting policy.
 - Experimental performance; benchmarks are local references only.

{tensorstudio-1.3.3 → tensorstudio-1.3.5}/README.md

@@ -7,7 +7,7 @@
 TensorStudio is a compact C++ tensor and autograd engine with a Python API for
 learning, experimentation, and lightweight ML workloads.
 
-TensorStudio `1.3.3` is a CPU-only stable API foundation. It is eager-only,
+TensorStudio `1.3.5` is a CPU-only stable API foundation. It is eager-only,
 intentionally small, and not a replacement for mature ML frameworks.
 
 ## Install
@@ -408,8 +408,10 @@ tokens or print secrets.
   pretrained model zoos, detection/segmentation training stacks, video IO, or
   GPU image kernels yet.
 - ONNX support is export-only for a limited set of TensorStudio modules.
-- Reductions support all-element or single-axis reductions, not tuple-axis
-  reductions yet.
+- Reductions support all-element, single-axis, and tuple/list-axis reductions
+  for `sum`, `mean`, `max`, and `min`.
+- Arg reductions support all-element flat indices or one axis at a time for
+  `argmax` and `argmin`.
 - No sparse tensors or advanced indexing.
 - Dtype casting is basic and does not include a full promotion/casting policy.
 - Experimental performance; benchmarks are local references only.

{tensorstudio-1.3.3 → tensorstudio-1.3.5}/docs/Autograde/index.md

@@ -40,7 +40,7 @@ Calling `backward()`:
 
 ## Supported Gradient Operations
 
-`1.3.3` supports gradients for:
+`1.3.5` supports gradients for:
 
 - `add`
 - `sub`

{tensorstudio-1.3.3 → tensorstudio-1.3.5}/docs/Benchmarks/index.md

@@ -55,7 +55,7 @@ equivalents where available.
 ## Interpreting Results
 
 Expect NumPy and PyTorch to be faster for many medium and large operations
-because they use highly optimized native kernels. TensorStudio `1.3.3` still
+because they use highly optimized native kernels. TensorStudio `1.3.5` still
 favors clarity, portability, and a compact C++ implementation over full kernel
 library performance.
 

{tensorstudio-1.3.3 → tensorstudio-1.3.5}/docs/Hardware/cpu-backend.md

@@ -1,6 +1,6 @@
 # CPU Backend
 
-TensorStudio `1.3.3` supports CPU tensors only.
+TensorStudio `1.3.5` supports CPU tensors only.
 
 ## Device Abstraction
 

{tensorstudio-1.3.3 → tensorstudio-1.3.5}/docs/Neural Networks/training.md

@@ -80,7 +80,7 @@ w = ts.randn((4, 4), seed=123)
 
 ## Practical Tips
 
-- Keep tensors modest in `1.3.3`; kernels are simple CPU loops.
+- Keep tensors modest in `1.3.5`; kernels are simple CPU loops.
 - Prefer `float32` unless you need `float64`.
 - Compare with NumPy in tests for expected numerical values.
 - Watch shape errors carefully around broadcasting and matrix multiplication.

{tensorstudio-1.3.3 → tensorstudio-1.3.5}/docs/Roadmap/index.md

@@ -1,6 +1,6 @@
 # Roadmap
 
-TensorStudio `1.3.3` is a CPU-first tensor, autograd, neural-network, vision,
+TensorStudio `1.3.5` is a CPU-first tensor, autograd, neural-network, vision,
 project, serialization, and ONNX-export foundation. The long-term direction is
 to become a strong compact ML framework for learning, experimentation, and
 lightweight workloads while staying honest about the scale of mature systems
@@ -21,8 +21,9 @@ tensor semantics.
 - Completed in `1.3.2`: add exhaustive broadcasting tests for every binary
   elementwise arithmetic and comparison operation, including functional and
   operator overload paths.
-- Add tuple-axis reductions for `sum`, `mean`, `max`, and `min`.
-- Add arg reductions: `argmax` and `argmin`.
+- Completed in `1.3.4`: add tuple-axis reductions for `sum`, `mean`,
+  `max`, and `min`.
+- Completed in `1.3.5`: add arg reductions: `argmax` and `argmin`.
 - Add `where`, `maximum`, `minimum`, `clip`, and richer comparison helpers.
 - Add full NumPy-style indexing and slicing for common cases.
 - Add more view/layout operations: `squeeze`, `unsqueeze`, `permute`, and
@@ -95,6 +96,9 @@ families.
 The project utilities should evolve into a clean small-project workflow without
 becoming an opaque training framework.
 
+- Add dataset creation helpers for images, tabular arrays, labels, and
+  TensorStudio tensors, with deterministic train/validation splitting and
+  metadata summaries.
 - Add callbacks for checkpointing, early stopping, CSV logging, and learning
   rate logging.
 - Add a metrics package for regression, classification, and multilabel tasks.
@@ -131,7 +135,18 @@ compatibility.
 - Expand NPZ metadata for richer model and optimizer state.
 - Add `safetensors` support for safe tensor weight storage.
 - Expand ONNX export coverage for more TensorStudio modules and ops.
+- Add ONNX metadata inspection for supported files, including graph inputs,
+  outputs, opset, initializer summaries, node counts, and model producer
+  fields.
 - Add ONNX import for a supported subset of static graphs.
+- Add execution support for imported ONNX graphs only where TensorStudio has
+  correct matching tensor ops; clearly reject unsupported operators and dynamic
+  graph features.
+- Add model metadata extraction for TensorStudio checkpoints, NPZ bundles, and
+  supported ONNX files. Do not advertise metadata support for formats that
+  cannot be parsed safely or accurately.
+- Add import/export research for practical neural-network model formats that
+  can be supported without unsafe execution or fake compatibility layers.
 - Add versioned checkpoint metadata and compatibility checks.
 - Add model card or metadata export for packaged examples.
 - Consider DLPack-style interop after tensor/device semantics are mature.

{tensorstudio-1.3.3 → tensorstudio-1.3.5}/docs/Usage/api.md

@@ -146,6 +146,16 @@ Math:
 - `mean(axis=None, keepdims=False)`
 - `max(axis=None, keepdims=False)`
 - `min(axis=None, keepdims=False)`
+- `argmax(axis=None, keepdims=False)`
+- `argmin(axis=None, keepdims=False)`
+
+For reductions, `axis` may be `None`, an int, or a tuple/list of ints.
+Negative axes are normalized against the input rank, duplicate axes are
+rejected, and `axis=()` is a no-op identity.
+
+For arg reductions, `axis` may be `None` or an int. All-element arg reductions
+return the flat index of the first selected value. Axis arg reductions return
+`int64` indices along that axis and are not differentiable.
 - `relu()`
 - `sigmoid()`
 - `tanh()`
@@ -295,6 +305,9 @@ Optimizers implement `zero_grad`, `step`, `state_dict`, and `load_state_dict`.
 - `std(input, axis=None, keepdims=False, correction=0)`
 - `norm(input, ord=2, axis=None, keepdims=False)`
 
+These helpers accept the same reduction axis forms as native reductions:
+`None`, an int, or a tuple/list of ints.
+
 These helpers compose TensorStudio tensors and keep autograd support.
 
 ## `tensorstudio.project`

{tensorstudio-1.3.3 → tensorstudio-1.3.5}/docs/Usage/interchange.md

@@ -1,6 +1,6 @@
 # ONNX Interchange
 
-TensorStudio `1.3.3` includes a limited ONNX exporter for supported
+TensorStudio `1.3.5` includes a limited ONNX exporter for supported
 TensorStudio module stacks.
 
 Install the optional dependency:

{tensorstudio-1.3.3 → tensorstudio-1.3.5}/docs/Usage/serialization.md

@@ -1,6 +1,6 @@
 # Serialization
 
-TensorStudio `1.3.3` provides two serialization paths: trusted pickle
+TensorStudio `1.3.5` provides two serialization paths: trusted pickle
 roundtrips for internal objects and non-pickle NPZ files for tensor/state_dict
 interchange.
 

{tensorstudio-1.3.3 → tensorstudio-1.3.5}/docs/Usage/tensors.md

@@ -184,15 +184,20 @@ Matrix and reductions:
 - `mean`
 - `max`
 - `min`
+- `argmax`
+- `argmin`
 
-Reductions operate over all elements by default and support one axis at a time
-with `keepdims`:
+Reductions operate over all elements by default and support an int axis or a
+tuple/list of axes with `keepdims`:
 
 ```python
 x = ts.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
 x.sum(axis=1)
 x.mean(axis=0, keepdims=True)
 x.max(axis=-1)
+x.reshape((1, 2, 3)).sum(axis=(0, 2))
+x.argmax()
+x.argmin(axis=1)
 ```
 
 Higher-level helpers in `tensorstudio.math` include variance, standard

{tensorstudio-1.3.3 → tensorstudio-1.3.5}/docs/index.md

@@ -3,7 +3,7 @@
 TensorStudio is a compact C++ tensor and autograd engine with a Python API for
 learning, experimentation, and lightweight ML workloads.
 
-`1.3.3` is a CPU-only stable API foundation for the tensor, autograd,
+`1.3.5` is a CPU-only stable API foundation for the tensor, autograd,
 neural-network, optimizer, data, project, serialization, ONNX export, vision,
 docs, packaging, and wheel workflows.
 
@@ -26,10 +26,11 @@ TensorStudio is:
   `empty`, `rand`, `randn`, `arange`, `eye`, `linspace`, and matching
   `*_like` helpers.
 - NumPy-style broadcasting for binary elementwise operations.
-- Arithmetic, comparisons, matrix multiplication, all-element and single-axis
-  reductions, common activations, trigonometric functions, inverse
-  trigonometric functions, `log1p`, `sqrt`, `rsqrt`, `abs`, `clamp`, CPU NCHW
-  `conv2d`, `max_pool2d`, `avg_pool2d`, reshape, flatten, and 2D transpose.
+- Arithmetic, comparisons, matrix multiplication, all-element, single-axis, and
+  tuple-axis reductions, arg reductions, common activations, trigonometric
+  functions, inverse trigonometric functions, `log1p`, `sqrt`, `rsqrt`, `abs`,
+  `clamp`, CPU NCHW `conv2d`, `max_pool2d`, `avg_pool2d`, reshape, flatten,
+  and 2D transpose.
 - Higher-level `tensorstudio.math` helpers for variance, standard deviation,
   norms, square, and reciprocal.
 - Basic dtype casting plus native `concat` and `stack`.
@@ -61,7 +62,7 @@ TensorStudio prioritizes:
 
 TensorStudio does not currently include CUDA, Metal, distributed execution,
 graph compilation, grouped convolution, adaptive/global pooling, advanced
-indexing, sparse tensors, mixed precision, tuple-axis reductions, a high
+indexing, sparse tensors, mixed precision, a high
 performance kernel library, pretrained vision model zoo, detection/segmentation
 training stack, video IO, an ONNX runtime, or ONNX import. Benchmarks are rough
 local references only; TensorStudio wins some small eager CPU cases but is not

{tensorstudio-1.3.3 → tensorstudio-1.3.5}/docs/publishing.md

@@ -22,8 +22,8 @@ Windows, Linux, and macOS.
 After the release checklist passes, tag the stable release:
 
 ```bash
-git tag v1.3.3
-git push origin v1.3.3
+git tag v1.3.5
+git push origin v1.3.5
 ```
 
 Stable release notes must avoid unsupported performance claims. Publish measured
@@ -63,7 +63,7 @@ clean environment:
 python -m venv .venv-testpypi
 . .venv-testpypi/bin/activate
 python -m pip install -U pip
-python -m pip install --index-url https://test.pypi.org/simple/ --extra-index-url https://pypi.org/simple/ tensorstudio==1.3.3
+python -m pip install --index-url https://test.pypi.org/simple/ --extra-index-url https://pypi.org/simple/ tensorstudio==1.3.5
 python -c "import tensorstudio as ts; import tensorstudio._C; print(ts.__version__)"
 deactivate
 ```
@@ -74,7 +74,7 @@ On Windows PowerShell:
 python -m venv .venv-testpypi
 .\.venv-testpypi\Scripts\Activate.ps1
 python -m pip install -U pip
-python -m pip install --index-url https://test.pypi.org/simple/ --extra-index-url https://pypi.org/simple/ tensorstudio==1.3.3
+python -m pip install --index-url https://test.pypi.org/simple/ --extra-index-url https://pypi.org/simple/ tensorstudio==1.3.5
 python -c "import tensorstudio as ts; import tensorstudio._C; print(ts.__version__)"
 deactivate
 ```

{tensorstudio-1.3.3 → tensorstudio-1.3.5}/docs/quickstart.md

@@ -109,13 +109,15 @@ print(ts.sin(x).tolist())
 print(ts.math.std(x).item())
 ```
 
-Reductions can operate over all elements or one axis:
+Reductions can operate over all elements, one axis, or a tuple/list of axes:
 
 ```python
 x = ts.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
 print(x.sum().item())
 print(x.mean(axis=0).tolist())
 print(x.max(axis=1, keepdims=True).tolist())
+print(ts.arange(24).reshape((2, 3, 4)).sum(axis=(0, 2)).tolist())
+print(x.argmax().item())
 ```
 
 Tensors can be cast, concatenated, and stacked:

{tensorstudio-1.3.3 → tensorstudio-1.3.5}/include/tensorstudio/ops.hpp

@@ -66,6 +66,10 @@ Tensor max(const Tensor& input);
 Tensor max(const Tensor& input, int64_t axis, bool keepdims = false);
 Tensor min(const Tensor& input);
 Tensor min(const Tensor& input, int64_t axis, bool keepdims = false);
+Tensor argmax(const Tensor& input, bool keepdims = false);
+Tensor argmax(const Tensor& input, int64_t axis, bool keepdims = false);
+Tensor argmin(const Tensor& input, bool keepdims = false);
+Tensor argmin(const Tensor& input, int64_t axis, bool keepdims = false);
 Tensor relu(const Tensor& input);
 Tensor sigmoid(const Tensor& input);
 Tensor tanh(const Tensor& input);

{tensorstudio-1.3.3 → tensorstudio-1.3.5}/include/tensorstudio/version.hpp

@@ -2,6 +2,6 @@
 
 namespace tensorstudio {
 
-inline constexpr const char* version = "1.3.3";
+inline constexpr const char* version = "1.3.5";
 
 }  // namespace tensorstudio

{tensorstudio-1.3.3 → tensorstudio-1.3.5}/pyproject.toml

@@ -7,7 +7,7 @@ build-backend = "scikit_build_core.build"
 
 [project]
 name = "tensorstudio"
-version = "1.3.3"
+version = "1.3.5"
 description = "TensorStudio is a compact C++ tensor and autograd engine with a Python API for learning, experimentation, and lightweight ML workloads."
 readme = "README.md"
 requires-python = ">=3.10"

{tensorstudio-1.3.3 → tensorstudio-1.3.5}/python/tensorstudio/_C.pyi

@@ -4,6 +4,8 @@ from typing import Any
 
 import numpy as np
 
+AxisLike = int | tuple[int, ...] | list[int] | None
+
 class Tensor:
     requires_grad: bool
     grad: Tensor | None
@@ -23,10 +25,12 @@ class Tensor:
     def reshape(self, *shape: Any) -> Tensor: ...
     def flatten(self) -> Tensor: ...
     def transpose(self) -> Tensor: ...
-    def sum(self, axis: int | None = None, keepdims: bool = False) -> Tensor: ...
-    def mean(self, axis: int | None = None, keepdims: bool = False) -> Tensor: ...
-    def max(self, axis: int | None = None, keepdims: bool = False) -> Tensor: ...
-    def min(self, axis: int | None = None, keepdims: bool = False) -> Tensor: ...
+    def sum(self, axis: AxisLike = None, keepdims: bool = False) -> Tensor: ...
+    def mean(self, axis: AxisLike = None, keepdims: bool = False) -> Tensor: ...
+    def max(self, axis: AxisLike = None, keepdims: bool = False) -> Tensor: ...
+    def min(self, axis: AxisLike = None, keepdims: bool = False) -> Tensor: ...
+    def argmax(self, axis: int | None = None, keepdims: bool = False) -> Tensor: ...
+    def argmin(self, axis: int | None = None, keepdims: bool = False) -> Tensor: ...
     def relu(self) -> Tensor: ...
     def sigmoid(self) -> Tensor: ...
     def tanh(self) -> Tensor: ...
@@ -140,10 +144,12 @@ def avg_pool2d(
     padding_w: int = 0,
     count_include_pad: bool = False,
 ) -> Tensor: ...
-def sum(input: Tensor, axis: int | None = None, keepdims: bool = False) -> Tensor: ...
-def mean(input: Tensor, axis: int | None = None, keepdims: bool = False) -> Tensor: ...
-def max(input: Tensor, axis: int | None = None, keepdims: bool = False) -> Tensor: ...
-def min(input: Tensor, axis: int | None = None, keepdims: bool = False) -> Tensor: ...
+def sum(input: Tensor, axis: AxisLike = None, keepdims: bool = False) -> Tensor: ...
+def mean(input: Tensor, axis: AxisLike = None, keepdims: bool = False) -> Tensor: ...
+def max(input: Tensor, axis: AxisLike = None, keepdims: bool = False) -> Tensor: ...
+def min(input: Tensor, axis: AxisLike = None, keepdims: bool = False) -> Tensor: ...
+def argmax(input: Tensor, axis: int | None = None, keepdims: bool = False) -> Tensor: ...
+def argmin(input: Tensor, axis: int | None = None, keepdims: bool = False) -> Tensor: ...
 def relu(input: Tensor) -> Tensor: ...
 def sigmoid(input: Tensor) -> Tensor: ...
 def tanh(input: Tensor) -> Tensor: ...

{tensorstudio-1.3.3 → tensorstudio-1.3.5}/python/tensorstudio/__init__.py

@@ -9,6 +9,8 @@ from .grad_mode import is_grad_enabled, no_grad, set_grad_enabled
 from .interchange import export_onnx
 from .ops import (
     acos,
+    argmax,
+    argmin,
     asin,
     astype,
     atan,
@@ -50,6 +52,8 @@ __all__ = [
     "Tensor",
     "__version__",
     "acos",
+    "argmax",
+    "argmin",
     "arange",
     "asin",
     "astype",

{tensorstudio-1.3.3 → tensorstudio-1.3.5}/python/tensorstudio/_version.py

@@ -1,5 +1,5 @@
 """Package version."""
 
-__version__ = "1.3.3"
+__version__ = "1.3.5"
 
 __all__ = ["__version__"]

{tensorstudio-1.3.3 → tensorstudio-1.3.5}/python/tensorstudio/math.py

@@ -7,6 +7,25 @@ from typing import Literal
 from .tensor import Tensor
 
 NormOrder = Literal[1, 2, "fro", "inf"]
+AxisLike = int | tuple[int, ...] | list[int] | None
+
+
+def _normalized_axes(axis: AxisLike, ndim: int) -> tuple[int, ...]:
+    if axis is None:
+        return tuple(range(ndim))
+    axes = (axis,) if isinstance(axis, int) else tuple(axis)
+
+    normalized: list[int] = []
+    seen: set[int] = set()
+    for value in axes:
+        normalized_axis = value if value >= 0 else ndim + value
+        if normalized_axis < 0 or normalized_axis >= ndim:
+            raise ValueError("axis is out of range")
+        if normalized_axis in seen:
+            raise ValueError("duplicate reduction axis")
+        seen.add(normalized_axis)
+        normalized.append(normalized_axis)
+    return tuple(normalized)
 
 
 def square(input: Tensor) -> Tensor:
@@ -23,25 +42,23 @@ def reciprocal(input: Tensor) -> Tensor:
 
 def variance(
     input: Tensor,
-    axis: int | None = None,
+    axis: AxisLike = None,
     keepdims: bool = False,
     correction: int = 0,
 ) -> Tensor:
-    """Compute variance over all elements or one axis.
+    """Compute variance over all elements, one axis, or a tuple/list of axes.
 
     ``correction=1`` gives the usual sample variance denominator. TensorStudio
-    currently supports a single optional axis, matching the C++ reduction API.
+    uses the product of the reduced axes as the denominator base.
     """
 
     if correction < 0:
         raise ValueError("correction must be non-negative")
-    if axis is None:
-        denom = input.size - correction
-    else:
-        normalized_axis = axis if axis >= 0 else input.ndim + axis
-        if normalized_axis < 0 or normalized_axis >= input.ndim:
-            raise ValueError("axis is out of range")
-        denom = input.shape[normalized_axis] - correction
+    axes = _normalized_axes(axis, input.ndim)
+    denom = 1
+    for normalized_axis in axes:
+        denom *= input.shape[normalized_axis]
+    denom -= correction
     if denom <= 0:
         raise ValueError("variance denominator must be positive")
     mean = input.mean(axis=axis, keepdims=True)
@@ -52,11 +69,11 @@ def variance(
 
 def std(
     input: Tensor,
-    axis: int | None = None,
+    axis: AxisLike = None,
     keepdims: bool = False,
     correction: int = 0,
 ) -> Tensor:
-    """Compute standard deviation over all elements or one axis."""
+    """Compute standard deviation over all elements, one axis, or a tuple/list of axes."""
 
     return variance(input, axis=axis, keepdims=keepdims, correction=correction).sqrt()
 
@@ -64,10 +81,10 @@ def std(
 def norm(
     input: Tensor,
     ord: NormOrder = 2,
-    axis: int | None = None,
+    axis: AxisLike = None,
     keepdims: bool = False,
 ) -> Tensor:
-    """Compute a vector-style norm over all elements or one axis."""
+    """Compute a vector-style norm over all elements, one axis, or a tuple/list of axes."""
 
     magnitude = input.abs()
     if ord == 1:
@@ -79,4 +96,12 @@ def norm(
     raise ValueError("ord must be 1, 2, 'fro', or 'inf'")
 
 
-__all__ = ["NormOrder", "norm", "reciprocal", "square", "std", "variance"]
+__all__ = [
+    "AxisLike",
+    "NormOrder",
+    "norm",
+    "reciprocal",
+    "square",
+    "std",
+    "variance",
+]

{tensorstudio-1.3.3 → tensorstudio-1.3.5}/python/tensorstudio/ops.py

@@ -8,6 +8,7 @@ from . import _C
 from .tensor import Tensor
 
 PairLike = int | tuple[int, int] | list[int]
+AxisLike = int | tuple[int, ...] | list[int] | None
 
 
 def _pair(value: PairLike, name: str) -> tuple[int, int]:
@@ -147,22 +148,30 @@ def avg_pool2d(
     )
 
 
-def sum(input: Tensor, axis: int | None = None, keepdims: bool = False) -> Tensor:  # noqa: A001
+def sum(input: Tensor, axis: AxisLike = None, keepdims: bool = False) -> Tensor:  # noqa: A001
     return _C.sum(input, axis, keepdims)
 
 
-def mean(input: Tensor, axis: int | None = None, keepdims: bool = False) -> Tensor:
+def mean(input: Tensor, axis: AxisLike = None, keepdims: bool = False) -> Tensor:
     return _C.mean(input, axis, keepdims)
 
 
-def max(input: Tensor, axis: int | None = None, keepdims: bool = False) -> Tensor:  # noqa: A001
+def max(input: Tensor, axis: AxisLike = None, keepdims: bool = False) -> Tensor:  # noqa: A001
     return _C.max(input, axis, keepdims)
 
 
-def min(input: Tensor, axis: int | None = None, keepdims: bool = False) -> Tensor:  # noqa: A001
+def min(input: Tensor, axis: AxisLike = None, keepdims: bool = False) -> Tensor:  # noqa: A001
     return _C.min(input, axis, keepdims)
 
 
+def argmax(input: Tensor, axis: int | None = None, keepdims: bool = False) -> Tensor:
+    return _C.argmax(input, axis, keepdims)
+
+
+def argmin(input: Tensor, axis: int | None = None, keepdims: bool = False) -> Tensor:
+    return _C.argmin(input, axis, keepdims)
+
+
 def relu(input: Tensor) -> Tensor:
     return _C.relu(input)
 
@@ -259,6 +268,8 @@ __all__ = [
     "add",
     "abs",
     "acos",
+    "argmax",
+    "argmin",
     "asin",
     "astype",
     "atan",

{tensorstudio-1.3.3 → tensorstudio-1.3.5}/src/bindings/bind_ops.cpp

@@ -1,5 +1,7 @@
 #include "bindings.hpp"
 
+#include <utility>
+
 #include <pybind11/pybind11.h>
 #include <pybind11/stl.h>
 
@@ -72,28 +74,46 @@ void bind_ops(py::module_& module) {
       py::arg("padding_w") = 0,
       py::arg("count_include_pad") = false);
   module.def("sum", [](const Tensor& input, py::object axis, bool keepdims) {
-    if (axis.is_none()) {
-      return sum(input);
-    }
-    return sum(input, py::cast<int64_t>(axis), keepdims);
+    return reduce_from_py(input, std::move(axis), keepdims, "sum", [](const Tensor& tensor) {
+      return sum(tensor);
+    }, [](const Tensor& tensor, int64_t normalized_axis, bool keep) {
+      return sum(tensor, normalized_axis, keep);
+    });
   }, py::arg("input"), py::arg("axis") = py::none(), py::arg("keepdims") = false);
   module.def("mean", [](const Tensor& input, py::object axis, bool keepdims) {
-    if (axis.is_none()) {
-      return mean(input);
-    }
-    return mean(input, py::cast<int64_t>(axis), keepdims);
+    return reduce_from_py(input, std::move(axis), keepdims, "mean", [](const Tensor& tensor) {
+      return mean(tensor);
+    }, [](const Tensor& tensor, int64_t normalized_axis, bool keep) {
+      return mean(tensor, normalized_axis, keep);
+    });
   }, py::arg("input"), py::arg("axis") = py::none(), py::arg("keepdims") = false);
   module.def("max", [](const Tensor& input, py::object axis, bool keepdims) {
-    if (axis.is_none()) {
-      return max(input);
-    }
-    return max(input, py::cast<int64_t>(axis), keepdims);
+    return reduce_from_py(input, std::move(axis), keepdims, "max", [](const Tensor& tensor) {
+      return max(tensor);
+    }, [](const Tensor& tensor, int64_t normalized_axis, bool keep) {
+      return max(tensor, normalized_axis, keep);
+    });
   }, py::arg("input"), py::arg("axis") = py::none(), py::arg("keepdims") = false);
   module.def("min", [](const Tensor& input, py::object axis, bool keepdims) {
-    if (axis.is_none()) {
-      return min(input);
-    }
-    return min(input, py::cast<int64_t>(axis), keepdims);
+    return reduce_from_py(input, std::move(axis), keepdims, "min", [](const Tensor& tensor) {
+      return min(tensor);
+    }, [](const Tensor& tensor, int64_t normalized_axis, bool keep) {
+      return min(tensor, normalized_axis, keep);
+    });
+  }, py::arg("input"), py::arg("axis") = py::none(), py::arg("keepdims") = false);
+  module.def("argmax", [](const Tensor& input, py::object axis, bool keepdims) {
+    return arg_reduce_from_py(input, std::move(axis), keepdims, "argmax", [](const Tensor& tensor, bool keep) {
+      return argmax(tensor, keep);
+    }, [](const Tensor& tensor, int64_t normalized_axis, bool keep) {
+      return argmax(tensor, normalized_axis, keep);
+    });
+  }, py::arg("input"), py::arg("axis") = py::none(), py::arg("keepdims") = false);
+  module.def("argmin", [](const Tensor& input, py::object axis, bool keepdims) {
+    return arg_reduce_from_py(input, std::move(axis), keepdims, "argmin", [](const Tensor& tensor, bool keep) {
+      return argmin(tensor, keep);
+    }, [](const Tensor& tensor, int64_t normalized_axis, bool keep) {
+      return argmin(tensor, normalized_axis, keep);
+    });
   }, py::arg("input"), py::arg("axis") = py::none(), py::arg("keepdims") = false);
   module.def("relu", &relu);
   module.def("sigmoid", &sigmoid);