Function go_with_map_spec

src/python/py_runtime.rs:459–495

()

    #[test]
    fn go_with_map_spec() {
        pyo3::prepare_freethreaded_python();
        Python::with_gil(|py| {
            use crate::buffer::inner::{Buffer, DType};
            use crate::ir::compiler::ArgValue;
            use crate::ir::expr::Expr;
            use crate::ir::kernel::{ArgSpec, KernelKind, KernelSpec, TensorSpec};
            use crate::python::py_kernel::{PyKernelSpec, PyMapSpec};
            use std::collections::HashMap;

            let inner_spec = KernelSpec {
                kind: KernelKind::Elementwise,
                args: vec![ArgSpec {
                    name: "x".to_string(),
                    dtype: DType::F64,
                    is_scalar: false,
                }],
                output: TensorSpec { dtype: DType::F64 },
                expr: Expr::ArgRef(0),
            };
            let spec = PyKernelSpec { inner: inner_spec };
            let mut args = HashMap::new();
            args.insert(
                "x".to_string(),
                ArgValue::Buffer(Buffer::from_f64_vec(vec![1.0, 2.0])),
            );
            let map_spec = PyMapSpec { spec, args };
            let py_map = map_spec.into_pyobject(py).unwrap();

            let runtime = PyRuntimeModule::new();
            let call_args = empty_args(py);
            let task = runtime
                .go(py, py_map.as_any(), &call_args, None, None)
                .unwrap();
            assert!(task.inner.is_done());
        });
    }

    #[test]
    fn go_with_reduce_spec() {