1use pyo3::prelude::*;
4use ronn_core::Tensor;
5
6#[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 fn shape(&self) -> Vec<usize> {
22 self.inner.shape()
23 }
24
25 fn dtype(&self) -> String {
27 format!("{:?}", self.inner.dtype())
28 }
29
30 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}