diff --git a/Cargo.toml b/Cargo.toml index e4664a34..a30530a7 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -25,8 +25,8 @@ plotly = { version = "0.13.5" } argmin = { version = "0.11.0" } argmin-math = { version = "0.5.1" } argmin-observer-slog = { version = "0.2.0" } -ort = "=2.0.0-rc.9" -ort-sys = { version = "=2.0.0-rc.9", default-features = false } +ort = "=2.0.0-rc.10" +ort-sys = { version = "=2.0.0-rc.10", default-features = false } [profile.profiling] inherits = "release" diff --git a/examples/neural-ode-weather-prediction/Cargo.toml b/examples/neural-ode-weather-prediction/Cargo.toml index 166d19c6..4da2ecd3 100644 --- a/examples/neural-ode-weather-prediction/Cargo.toml +++ b/examples/neural-ode-weather-prediction/Cargo.toml @@ -20,4 +20,4 @@ ort = { workspace = true, optional = true } ort-sys = { workspace = true, optional = true } ndarray = "0.16.1" csv = "1.3.1" -rand = "0.9.0" \ No newline at end of file +rand = "0.9.0" diff --git a/examples/neural-ode-weather-prediction/src/main.rs b/examples/neural-ode-weather-prediction/src/main.rs index ff61a4f9..e9fbdb13 100644 --- a/examples/neural-ode-weather-prediction/src/main.rs +++ b/examples/neural-ode-weather-prediction/src/main.rs @@ -15,6 +15,7 @@ use ndarray::Array1; use ort::error::Result; use ort::inputs; use ort::session::{builder::GraphOptimizationLevel, Session}; +use ort::value::{TensorRef, ValueType}; use plotly::common::{DashType, Line, Mode}; use plotly::layout::{Axis, GridPattern, LayoutGrid}; use plotly::{Layout, Plot, Scatter}; @@ -31,10 +32,10 @@ const BASE_DATA_DIR: &str = "examples/neural-ode-weather-prediction/src/data/"; const BASE_OUTPUT_DIR: &str = "examples/neural-ode-weather-prediction/"; struct NeuralOde { - rhs: Session, - rhs_jac_mul: Session, - rhs_jac_transpose_mul: Session, - rhs_sens_transpose_mul: Session, + rhs: RefCell, + rhs_jac_mul: RefCell, + rhs_jac_transpose_mul: RefCell, + rhs_sens_transpose_mul: RefCell, input_y: RefCell>, input_v: RefCell>, input_p: Array1, @@ -58,8 +59,17 @@ impl NeuralOde { let mut nparams = 0; for input in rhs.inputs.iter() { if input.name == "p" { - nparams = input.input_type.tensor_dimensions().unwrap()[0] as usize; - break; + if let ValueType::Tensor { shape, .. } = &input.input_type { + let dim = shape + .get(0) + .copied() + .expect("p input should have at least one dimension"); + if dim < 0 { + panic!("p input has dynamic dimension; cannot infer parameter count"); + } + nparams = dim as usize; + break; + } } } let mut rng = rand::rng(); @@ -69,10 +79,10 @@ impl NeuralOde { Ok(Self { y0, - rhs, - rhs_jac_mul, - rhs_jac_transpose_mul, - rhs_sens_transpose_mul, + rhs: RefCell::new(rhs), + rhs_jac_mul: RefCell::new(rhs_jac_mul), + rhs_jac_transpose_mul: RefCell::new(rhs_jac_transpose_mul), + rhs_sens_transpose_mul: RefCell::new(rhs_sens_transpose_mul), input_p: params, input_v: RefCell::new(y0_ndarray.clone()), input_y: RefCell::new(y0_ndarray), @@ -200,21 +210,19 @@ impl NonLinearOp for Rhs<'_> { .iter_mut() .zip(x.inner().iter()) .for_each(|(y, x)| *y = *x as f32); - let outputs = self - .0 - .rhs - .run( - inputs![ - "p" => self.0.input_p.view(), - "y" => y_input.view(), - ] - .unwrap(), - ) + let p_slice = self.0.input_p.as_slice().unwrap(); + let y_slice = y_input.as_slice().unwrap(); + let mut rhs = self.0.rhs.borrow_mut(); + let outputs = rhs + .run(inputs![ + "p" => TensorRef::from_array_view(([p_slice.len()], p_slice)).unwrap(), + "y" => TensorRef::from_array_view(([y_slice.len()], y_slice)).unwrap(), + ]) .unwrap(); - let y_data = outputs["Identity_1:0"].try_extract_tensor::().unwrap(); + let (_shape, y_data) = outputs["Identity_1:0"].try_extract_tensor::().unwrap(); y.inner_mut() .iter_mut() - .zip(y_data.as_slice().unwrap()) + .zip(y_data.iter()) .for_each(|(y, x)| *y = *x as f64); } } @@ -231,22 +239,21 @@ impl NonLinearOpJacobian for Rhs<'_> { .iter_mut() .zip(v.inner().iter()) .for_each(|(v, x)| *v = *x as f32); - let outputs = self - .0 - .rhs_jac_mul - .run( - inputs![ - "y" => y_input.view(), - "v" => v_input.view(), - "p" => self.0.input_p.view(), - ] - .unwrap(), - ) + let p_slice = self.0.input_p.as_slice().unwrap(); + let y_slice = y_input.as_slice().unwrap(); + let v_slice = v_input.as_slice().unwrap(); + let mut rhs_jac_mul = self.0.rhs_jac_mul.borrow_mut(); + let outputs = rhs_jac_mul + .run(inputs![ + "y" => TensorRef::from_array_view(([y_slice.len()], y_slice)).unwrap(), + "v" => TensorRef::from_array_view(([v_slice.len()], v_slice)).unwrap(), + "p" => TensorRef::from_array_view(([p_slice.len()], p_slice)).unwrap(), + ]) .unwrap(); - let y_data = outputs["Identity_1:0"].try_extract_tensor::().unwrap(); + let (_shape, y_data) = outputs["Identity_1:0"].try_extract_tensor::().unwrap(); y.inner_mut() .iter_mut() - .zip(y_data.as_slice().unwrap()) + .zip(y_data.iter()) .for_each(|(y, x)| *y = *x as f64); } } @@ -263,22 +270,21 @@ impl NonLinearOpAdjoint for Rhs<'_> { .iter_mut() .zip(v.inner().iter()) .for_each(|(v, x)| *v = *x as f32); - let outputs = self - .0 - .rhs_jac_transpose_mul - .run( - inputs![ - "y" => y_input.view(), - "v" => v_input.view(), - "p" => self.0.input_p.view(), - ] - .unwrap(), - ) + let p_slice = self.0.input_p.as_slice().unwrap(); + let y_slice = y_input.as_slice().unwrap(); + let v_slice = v_input.as_slice().unwrap(); + let mut rhs_jac_transpose_mul = self.0.rhs_jac_transpose_mul.borrow_mut(); + let outputs = rhs_jac_transpose_mul + .run(inputs![ + "y" => TensorRef::from_array_view(([y_slice.len()], y_slice)).unwrap(), + "v" => TensorRef::from_array_view(([v_slice.len()], v_slice)).unwrap(), + "p" => TensorRef::from_array_view(([p_slice.len()], p_slice)).unwrap(), + ]) .unwrap(); - let y_data = outputs["Identity_1:0"].try_extract_tensor::().unwrap(); + let (_shape, y_data) = outputs["Identity_1:0"].try_extract_tensor::().unwrap(); y.inner_mut() .iter_mut() - .zip(y_data.as_slice().unwrap()) + .zip(y_data.iter()) .for_each(|(y, x)| *y = *x as f64); } } @@ -295,22 +301,21 @@ impl NonLinearOpSensAdjoint for Rhs<'_> { .iter_mut() .zip(v.inner().iter()) .for_each(|(v, x)| *v = *x as f32); - let outputs = self - .0 - .rhs_sens_transpose_mul - .run( - inputs![ - "y" => y_input.view(), - "v" => v_input.view(), - "p" => self.0.input_p.view(), - ] - .unwrap(), - ) + let p_slice = self.0.input_p.as_slice().unwrap(); + let y_slice = y_input.as_slice().unwrap(); + let v_slice = v_input.as_slice().unwrap(); + let mut rhs_sens_transpose_mul = self.0.rhs_sens_transpose_mul.borrow_mut(); + let outputs = rhs_sens_transpose_mul + .run(inputs![ + "y" => TensorRef::from_array_view(([y_slice.len()], y_slice)).unwrap(), + "v" => TensorRef::from_array_view(([v_slice.len()], v_slice)).unwrap(), + "p" => TensorRef::from_array_view(([p_slice.len()], p_slice)).unwrap(), + ]) .unwrap(); - let y_data = outputs["Identity_1:0"].try_extract_tensor::().unwrap(); + let (_shape, y_data) = outputs["Identity_1:0"].try_extract_tensor::().unwrap(); y.inner_mut() .iter_mut() - .zip(y_data.as_slice().unwrap()) + .zip(y_data.iter()) .for_each(|(y, x)| *y = *x as f64); } }