Skip to main content

ronn/
tensor.rs

1//! Tensor class for Python bindings
2
3use pyo3::prelude::*;
4use ronn_core::Tensor;
5
6/// Tensor wrapper for Python
7#[pyclass(name = "Tensor")]
8pub struct PyTensor {
9    pub(crate) inner: Tensor,
10}
11
12impl PyTensor {
13    pub fn new(tensor: Tensor) -> Self {
14        Self { inner: tensor }
15    }
16}
17
18#[pymethods]
19impl PyTensor {
20    /// Get tensor shape
21    fn shape(&self) -> Vec<usize> {
22        self.inner.shape()
23    }
24
25    /// Get tensor data type
26    fn dtype(&self) -> String {
27        format!("{:?}", self.inner.dtype())
28    }
29
30    /// Get number of elements
31    fn numel(&self) -> usize {
32        self.inner.numel()
33    }
34
35    fn __repr__(&self) -> String {
36        format!(
37            "Tensor(shape={:?}, dtype={}, numel={})",
38            self.shape(),
39            self.dtype(),
40            self.numel()
41        )
42    }
43}