1#![allow(dead_code)]
2#![allow(non_local_definitions)]
3use pyo3::prelude::*;
35
36mod error;
37mod model;
38mod session;
39mod tensor;
40
41pub use error::RonnError;
42pub use model::PyModel;
43pub use session::PySession;
44pub use tensor::PyTensor;
45
46#[pymodule]
48fn ronn(_py: Python, m: &PyModule) -> PyResult<()> {
49 m.add_class::<PyModel>()?;
50 m.add_class::<PySession>()?;
51 m.add_class::<PyTensor>()?;
52 m.add_class::<PyOptimizationLevel>()?;
53 m.add_class::<PyProviderType>()?;
54 m.add_class::<PyBatchConfig>()?;
55
56 m.add("__version__", env!("CARGO_PKG_VERSION"))?;
58
59 Ok(())
60}
61
62#[pyclass(name = "OptimizationLevel")]
64#[derive(Clone)]
65pub struct PyOptimizationLevel {
66 inner: ronn_core::OptimizationLevel,
67}
68
69#[pymethods]
70impl PyOptimizationLevel {
71 #[staticmethod]
73 fn none() -> Self {
74 Self {
75 inner: ronn_core::OptimizationLevel::None,
76 }
77 }
78
79 #[staticmethod]
81 fn basic() -> Self {
82 Self {
83 inner: ronn_core::OptimizationLevel::Basic,
84 }
85 }
86
87 #[staticmethod]
89 fn default() -> Self {
90 Self {
91 inner: ronn_core::OptimizationLevel::Basic,
92 }
93 }
94
95 #[staticmethod]
97 fn aggressive() -> Self {
98 Self {
99 inner: ronn_core::OptimizationLevel::Aggressive,
100 }
101 }
102}
103
104impl Default for PyOptimizationLevel {
105 fn default() -> Self {
106 Self::default()
107 }
108}
109
110#[pyclass(name = "ProviderType")]
112#[derive(Clone)]
113pub struct PyProviderType {
114 inner: ronn_core::ProviderId,
115}
116
117#[pymethods]
118impl PyProviderType {
119 #[staticmethod]
121 fn cpu() -> Self {
122 Self {
123 inner: ronn_core::ProviderId::CPU,
124 }
125 }
126
127 #[staticmethod]
129 fn gpu() -> Self {
130 Self {
131 inner: ronn_core::ProviderId::GPU,
132 }
133 }
134
135 #[staticmethod]
137 fn bitnet() -> Self {
138 Self {
139 inner: ronn_core::ProviderId::BitNet,
140 }
141 }
142
143 #[staticmethod]
145 fn wasm() -> Self {
146 Self {
147 inner: ronn_core::ProviderId::WebAssembly,
148 }
149 }
150}
151
152impl Default for PyProviderType {
153 fn default() -> Self {
154 Self::cpu()
155 }
156}
157
158#[pyclass(name = "BatchConfig")]
160#[derive(Clone)]
161pub struct PyBatchConfig {
162 #[pyo3(get, set)]
163 pub max_batch_size: usize,
165
166 #[pyo3(get, set)]
167 pub timeout_ms: u64,
169
170 #[pyo3(get, set)]
171 pub queue_capacity: usize,
173}
174
175#[pymethods]
176impl PyBatchConfig {
177 #[new]
178 #[pyo3(signature = (max_batch_size=32, timeout_ms=10, queue_capacity=1024))]
179 fn new(max_batch_size: usize, timeout_ms: u64, queue_capacity: usize) -> Self {
180 Self {
181 max_batch_size,
182 timeout_ms,
183 queue_capacity,
184 }
185 }
186}
187
188impl Default for PyBatchConfig {
189 fn default() -> Self {
190 Self::new(32, 10, 1024)
191 }
192}