tensorstudio 1.3.3__tar.gz → 1.3.4__tar.gz

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

@@ -2,6 +2,18 @@
 
 ## Unreleased
 
+## 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.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: tensorstudio
-Version: 1.3.3
+Version: 1.3.4
 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.4` 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,8 @@ 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`.
 - 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.4}/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.4` 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,8 @@ 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`.
 - 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.4}/docs/Autograde/index.md

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

{tensorstudio-1.3.3 → tensorstudio-1.3.4}/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.4` still
 favors clarity, portability, and a compact C++ implementation over full kernel
 library performance.
 

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

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

{tensorstudio-1.3.3 → tensorstudio-1.3.4}/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.4`; 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.4}/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.4` 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,7 +21,8 @@ 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`.
+- Completed in `1.3.4`: add tuple-axis reductions for `sum`, `mean`,
+  `max`, and `min`.
 - 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.
@@ -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.4}/docs/Usage/api.md

@@ -146,6 +146,10 @@ Math:
 - `mean(axis=None, keepdims=False)`
 - `max(axis=None, keepdims=False)`
 - `min(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.
 - `relu()`
 - `sigmoid()`
 - `tanh()`
@@ -295,6 +299,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.4}/docs/Usage/interchange.md

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

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

@@ -1,6 +1,6 @@
 # Serialization
 
-TensorStudio `1.3.3` provides two serialization paths: trusted pickle
+TensorStudio `1.3.4` 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.4}/docs/Usage/tensors.md

@@ -185,14 +185,15 @@ Matrix and reductions:
 - `max`
 - `min`
 
-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))
 ```
 
 Higher-level helpers in `tensorstudio.math` include variance, standard

{tensorstudio-1.3.3 → tensorstudio-1.3.4}/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.4` 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,8 +26,8 @@ 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
+- Arithmetic, comparisons, matrix multiplication, all-element, single-axis, and
+  tuple-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.
 - Higher-level `tensorstudio.math` helpers for variance, standard deviation,
@@ -61,7 +61,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.4}/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.4
+git push origin v1.3.4
 ```
 
 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.4
 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.4
 python -c "import tensorstudio as ts; import tensorstudio._C; print(ts.__version__)"
 deactivate
 ```

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

@@ -109,13 +109,14 @@ 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())
 ```
 
 Tensors can be cast, concatenated, and stacked:

{tensorstudio-1.3.3 → tensorstudio-1.3.4}/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.4";
 
 }  // namespace tensorstudio

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

@@ -7,7 +7,7 @@ build-backend = "scikit_build_core.build"
 
 [project]
 name = "tensorstudio"
-version = "1.3.3"
+version = "1.3.4"
 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.4}/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,10 @@ 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 relu(self) -> Tensor: ...
     def sigmoid(self) -> Tensor: ...
     def tanh(self) -> Tensor: ...
@@ -140,10 +142,10 @@ 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 relu(input: Tensor) -> Tensor: ...
 def sigmoid(input: Tensor) -> Tensor: ...
 def tanh(input: Tensor) -> Tensor: ...

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

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

{tensorstudio-1.3.3 → tensorstudio-1.3.4}/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.4}/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,19 +148,19 @@ 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)
 
 

{tensorstudio-1.3.3 → tensorstudio-1.3.4}/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,32 @@ 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("relu", &relu);
   module.def("sigmoid", &sigmoid);