tensorstudio 1.3.4__tar.gz → 1.3.5__tar.gz

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

@@ -2,6 +2,16 @@
 
 ## 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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: tensorstudio
-Version: 1.3.4
+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.4` 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
@@ -474,6 +474,8 @@ tokens or print secrets.
 - ONNX support is export-only for a limited set of TensorStudio modules.
 - 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.4 → 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.4` 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
@@ -410,6 +410,8 @@ tokens or print secrets.
 - ONNX support is export-only for a limited set of TensorStudio modules.
 - 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.4 → tensorstudio-1.3.5}/docs/Autograde/index.md

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

{tensorstudio-1.3.4 → 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.4` 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.4 → tensorstudio-1.3.5}/docs/Hardware/cpu-backend.md

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

{tensorstudio-1.3.4 → 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.4`; 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.4 → tensorstudio-1.3.5}/docs/Roadmap/index.md

@@ -1,6 +1,6 @@
 # Roadmap
 
-TensorStudio `1.3.4` 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
@@ -23,7 +23,7 @@ tensor semantics.
   operator overload paths.
 - Completed in `1.3.4`: add tuple-axis reductions for `sum`, `mean`,
   `max`, and `min`.
-- Add arg reductions: `argmax` and `argmin`.
+- 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

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

@@ -146,10 +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()`

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

@@ -1,6 +1,6 @@
 # ONNX Interchange
 
-TensorStudio `1.3.4` 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.4 → tensorstudio-1.3.5}/docs/Usage/serialization.md

@@ -1,6 +1,6 @@
 # Serialization
 
-TensorStudio `1.3.4` 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.4 → tensorstudio-1.3.5}/docs/Usage/tensors.md

@@ -184,6 +184,8 @@ Matrix and reductions:
 - `mean`
 - `max`
 - `min`
+- `argmax`
+- `argmin`
 
 Reductions operate over all elements by default and support an int axis or a
 tuple/list of axes with `keepdims`:
@@ -194,6 +196,8 @@ 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.4 → 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.4` 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.
 
@@ -27,9 +27,10 @@ TensorStudio is:
   `*_like` helpers.
 - NumPy-style broadcasting for binary elementwise operations.
 - 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.
+  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`.

{tensorstudio-1.3.4 → 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.4
-git push origin v1.3.4
+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.4
+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.4
+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.4 → tensorstudio-1.3.5}/docs/quickstart.md

@@ -117,6 +117,7 @@ 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.4 → 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.4 → tensorstudio-1.3.5}/include/tensorstudio/version.hpp

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

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

@@ -7,7 +7,7 @@ build-backend = "scikit_build_core.build"
 
 [project]
 name = "tensorstudio"
-version = "1.3.4"
+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.4 → tensorstudio-1.3.5}/python/tensorstudio/_C.pyi

@@ -29,6 +29,8 @@ class 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: ...
@@ -146,6 +148,8 @@ 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.4 → 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.4 → tensorstudio-1.3.5}/python/tensorstudio/_version.py

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

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

@@ -164,6 +164,14 @@ def min(input: Tensor, axis: AxisLike = None, keepdims: bool = False) -> Tensor:
     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)
 
@@ -260,6 +268,8 @@ __all__ = [
     "add",
     "abs",
     "acos",
+    "argmax",
+    "argmin",
     "asin",
     "astype",
     "atan",

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

@@ -101,6 +101,20 @@ void bind_ops(py::module_& module) {
       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);
   module.def("tanh", &tanh);

{tensorstudio-1.3.4 → tensorstudio-1.3.5}/src/bindings/bind_tensor.cpp

@@ -305,6 +305,22 @@ Tensor reduce_from_py(
   return result;
 }
 
+Tensor arg_reduce_from_py(
+    const Tensor& input,
+    py::object axis,
+    bool keepdims,
+    const std::string& op_name,
+    const std::function<Tensor(const Tensor&, bool)>& reduce_all,
+    const std::function<Tensor(const Tensor&, int64_t, bool)>& reduce_axis) {
+  if (axis.is_none()) {
+    return reduce_all(input, keepdims);
+  }
+  if (!py::isinstance<py::int_>(axis)) {
+    throw ShapeError(op_name + " axis must be None or an int");
+  }
+  return reduce_axis(input, normalize_reduction_axis(input, py::cast<int64_t>(axis), op_name), keepdims);
+}
+
 py::array tensor_to_numpy(const Tensor& tensor) {
   py::array array(numpy_dtype(tensor.dtype()), ssize_shape(tensor.shape()));
   if (tensor.is_contiguous()) {
@@ -398,6 +414,20 @@ void bind_tensor(py::module_& module) {
           return min(input, normalized_axis, keep);
         });
       }, py::arg("axis") = py::none(), py::arg("keepdims") = false)
+      .def("argmax", [](const Tensor& self, py::object axis, bool keepdims) {
+        return arg_reduce_from_py(self, std::move(axis), keepdims, "argmax", [](const Tensor& input, bool keep) {
+          return argmax(input, keep);
+        }, [](const Tensor& input, int64_t normalized_axis, bool keep) {
+          return argmax(input, normalized_axis, keep);
+        });
+      }, py::arg("axis") = py::none(), py::arg("keepdims") = false)
+      .def("argmin", [](const Tensor& self, py::object axis, bool keepdims) {
+        return arg_reduce_from_py(self, std::move(axis), keepdims, "argmin", [](const Tensor& input, bool keep) {
+          return argmin(input, keep);
+        }, [](const Tensor& input, int64_t normalized_axis, bool keep) {
+          return argmin(input, normalized_axis, keep);
+        });
+      }, py::arg("axis") = py::none(), py::arg("keepdims") = false)
       .def("relu", &relu)
       .def("sigmoid", &sigmoid)
       .def("tanh", &tanh)

{tensorstudio-1.3.4 → tensorstudio-1.3.5}/src/bindings/bindings.hpp

@@ -27,6 +27,13 @@ Tensor reduce_from_py(
     const std::string& op_name,
     const std::function<Tensor(const Tensor&)>& reduce_all,
     const std::function<Tensor(const Tensor&, int64_t, bool)>& reduce_axis);
+Tensor arg_reduce_from_py(
+    const Tensor& input,
+    py::object axis,
+    bool keepdims,
+    const std::string& op_name,
+    const std::function<Tensor(const Tensor&, bool)>& reduce_all,
+    const std::function<Tensor(const Tensor&, int64_t, bool)>& reduce_axis);
 
 void bind_tensor(py::module_& module);
 void bind_ops(py::module_& module);

{tensorstudio-1.3.4 → tensorstudio-1.3.5}/src/core/ops.cpp

@@ -1922,6 +1922,59 @@ Tensor min_impl(const Tensor& input, bool track_grad) {
   return out;
 }
 
+Tensor arg_extreme_impl(const Tensor& input, bool keepdims, bool is_max, const std::string& op_name) {
+  if (input.numel() == 0) {
+    throw ShapeError(op_name + " is undefined for empty tensors");
+  }
+  Shape out_shape;
+  if (keepdims) {
+    out_shape.assign(input.shape().size(), 1);
+  }
+  Tensor out(out_shape, DType::Int64, false);
+  int64_t best_index = 0;
+  double best_value = input.value_at_logical(0);
+  for (int64_t i = 1; i < input.numel(); ++i) {
+    const double value = input.value_at_logical(i);
+    if (is_max ? value > best_value : value < best_value) {
+      best_value = value;
+      best_index = i;
+    }
+  }
+  out.set_value_at_logical(0, static_cast<double>(best_index));
+  return out;
+}
+
+Tensor arg_extreme_axis_impl(
+    const Tensor& input,
+    int64_t axis,
+    bool keepdims,
+    bool is_max,
+    const std::string& op_name) {
+  const int64_t normalized_axis = normalize_axis(axis, input.shape(), op_name);
+  if (input.shape()[static_cast<std::size_t>(normalized_axis)] == 0) {
+    throw ShapeError(op_name + " is undefined for empty reduction axes");
+  }
+  const Shape out_shape = reduction_output_shape(input.shape(), normalized_axis, keepdims);
+  Tensor out(out_shape, DType::Int64, false);
+  std::vector<int64_t> best_indices(static_cast<std::size_t>(out.numel()), -1);
+  std::vector<double> best_values(static_cast<std::size_t>(out.numel()), 0.0);
+  for (int64_t i = 0; i < input.numel(); ++i) {
+    const Shape input_coords = coordinates_from_linear(i, input.shape());
+    const int64_t out_linear =
+        reduction_output_linear(input_coords, input.shape(), out_shape, normalized_axis, keepdims);
+    const auto out_index = static_cast<std::size_t>(out_linear);
+    const double value = input.value_at_logical(i);
+    if (best_indices[out_index] == -1 || (is_max ? value > best_values[out_index] : value < best_values[out_index])) {
+      best_values[out_index] = value;
+      best_indices[out_index] = input_coords[static_cast<std::size_t>(normalized_axis)];
+    }
+  }
+  for (int64_t i = 0; i < out.numel(); ++i) {
+    out.set_value_at_logical(i, static_cast<double>(best_indices[static_cast<std::size_t>(i)]));
+  }
+  return out;
+}
+
 Tensor relu_impl(const Tensor& input, bool track_grad) {
   Tensor out = elementwise_unary_impl(input, input.dtype(), [](double value) { return std::max(0.0, value); });
   if (track_grad) {
@@ -2378,6 +2431,22 @@ Tensor min(const Tensor& input, int64_t axis, bool keepdims) {
   return min_axis_impl(input, axis, keepdims, true);
 }
 
+Tensor argmax(const Tensor& input, bool keepdims) {
+  return arg_extreme_impl(input, keepdims, true, "argmax");
+}
+
+Tensor argmax(const Tensor& input, int64_t axis, bool keepdims) {
+  return arg_extreme_axis_impl(input, axis, keepdims, true, "argmax");
+}
+
+Tensor argmin(const Tensor& input, bool keepdims) {
+  return arg_extreme_impl(input, keepdims, false, "argmin");
+}
+
+Tensor argmin(const Tensor& input, int64_t axis, bool keepdims) {
+  return arg_extreme_axis_impl(input, axis, keepdims, false, "argmin");
+}
+
 Tensor relu(const Tensor& input) {
   return relu_impl(input, true);
 }

{tensorstudio-1.3.4 → tensorstudio-1.3.5}/tests/test_import.py

@@ -4,7 +4,7 @@ from __future__ import annotations
 def test_import_and_version() -> None:
     import tensorstudio as ts
 
-    assert ts.__version__ == "1.3.4"
+    assert ts.__version__ == "1.3.5"
     assert hasattr(ts, "Tensor")
     assert hasattr(ts, "data")
     assert hasattr(ts, "nn")

{tensorstudio-1.3.4 → tensorstudio-1.3.5}/tests/test_ops.py

@@ -271,6 +271,48 @@ def test_tuple_axis_reductions_match_numpy() -> None:
         x.max(axis=(0, "bad"))  # type: ignore[arg-type]
 
 
+def test_arg_reductions_match_numpy() -> None:
+    values = np.array(
+        [
+            [[1.0, 4.0, 4.0], [7.0, -2.0, 5.0]],
+            [[3.0, 3.0, 9.0], [0.0, -5.0, 9.0]],
+        ],
+        dtype=np.float32,
+    )
+    x = ts.from_numpy(values)
+
+    assert x.argmax().dtype == "int64"
+    assert x.argmin().dtype == "int64"
+    assert x.argmax().item() == int(np.argmax(values))
+    assert x.argmin().item() == int(np.argmin(values))
+    np.testing.assert_array_equal(x.argmax(axis=0).numpy(), np.argmax(values, axis=0))
+    np.testing.assert_array_equal(x.argmin(axis=1).numpy(), np.argmin(values, axis=1))
+    np.testing.assert_array_equal(
+        x.argmax(axis=-1, keepdims=True).numpy(),
+        np.argmax(values, axis=-1, keepdims=True),
+    )
+    np.testing.assert_array_equal(
+        ts.ops.argmin(x, axis=2, keepdims=True).numpy(),
+        np.argmin(values, axis=2, keepdims=True),
+    )
+    np.testing.assert_array_equal(
+        ts.argmax(x, keepdims=True).numpy(),
+        np.argmax(values, keepdims=True),
+    )
+
+    ties = ts.tensor([2.0, 5.0, 5.0, 1.0])
+    assert ties.argmax().item() == 1
+
+    with pytest.raises(Exception, match="out of range"):
+        x.argmax(axis=3)
+    with pytest.raises(Exception, match="axis must be None or an int"):
+        x.argmin(axis=(0, 1))  # type: ignore[arg-type]
+    with pytest.raises(Exception, match="empty"):
+        ts.tensor([], dtype="float32").argmax()
+    with pytest.raises(Exception, match="empty reduction axes"):
+        ts.empty((2, 0)).argmin(axis=1)
+
+
 def test_cast_concat_and_stack_match_numpy() -> None:
     x = ts.tensor([1.2, 2.8, -3.1])
     as_int = x.astype("int32")