Method elementwise

src/python/py_kernel.rs:61–107

Create an elementwise kernel from an expression. Performs dense reindexing: `temp_ids` are mapped to 0, 1, 2, ... in sorted order.

(&self, expr: PyExpr)

    /// Create an elementwise kernel from an expression.
    /// Performs dense reindexing: `temp_ids` are mapped to 0, 1, 2, ... in sorted order.
    fn elementwise(&self, expr: PyExpr) -> PyResult<PyKernelSpec> {
        // Collect temp_ids referenced in the expression, sorted
        let referenced_temp_ids = expr.inner.referenced_args();
        let mut seen_names = HashSet::new();

        // Build dense mapping: temp_id -> dense_index
        let mapping: HashMap<usize, usize> = referenced_temp_ids
            .iter()
            .enumerate()
            .map(|(dense_idx, &temp_id)| (temp_id, dense_idx))
            .collect();

        // Rewrite ArgRefs to dense indices
        let rewritten_expr = expr.inner.rewrite_arg_refs(&mapping);

        // Build ArgSpecs with names from temp_args
        let arg_specs: Vec<ArgSpec> = referenced_temp_ids
            .iter()
            .map(|temp_id| -> PyResult<ArgSpec> {
                let name = expr.temp_args.get(temp_id).cloned().ok_or_else(|| {
                    PyErr::new::<pyo3::exceptions::PyValueError, _>(format!(
                        "missing temp arg metadata for arg index {temp_id}"
                    ))
                })?;

                if !seen_names.insert(name.clone()) {
                    return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(format!(
                        "duplicate arg name: {name}"
                    )));
                }

                Ok(ArgSpec {
                    name,
                    dtype: DType::F64,
                    is_scalar: false,
                })
            })
            .collect::<PyResult<Vec<_>>>()?;

        let inner = KernelSpec {
            kind: KernelKind::Elementwise,
            args: arg_specs,
            output: crate::ir::kernel::TensorSpec { dtype: DType::F64 },
            expr: rewritten_expr,
        };
        Ok(PyKernelSpec { inner })
    }

    /// Map a kernel spec over named arguments, returning a `MapSpec` for `rt.go()`.
    #[pyo3(signature = (spec, **kwargs))]