nuvec_core/
graph.rs

1//! The graph-level API: primitives, fused and custom nodes, and access to an array's node.
2//!
3//! This is what graph transformations such as `nuvec-compile` build on; everyday code
4//! never needs it. It is not covered by the API's stability pins, and its functions return
5//! `Result` (they are not array operations: recording a primitive validates it, see
6//! [`from_primitive`]).
7
8use crate::array::State;
9use crate::{Array, DType, Result};
10
11pub use crate::custom::CustomCall;
12pub use crate::fast::FastOp;
13pub use crate::fused::{Fused, FusedRef};
14pub use crate::primitive::{
15    ArgReduceOp, BinaryOp, Primitive, ReduceOp, ScanOp, ScatterReduce, UnaryOp,
16};
17pub use crate::shape::Flags;
18pub use crate::view::SliceArg;
19
20/// Records `primitive` applied to `inputs`, with the same validation as the operations,
21/// and returns one array per output. Failure is sticky: if an input has failed, the
22/// outputs fail with its error.
23pub fn from_primitive(primitive: Primitive, inputs: Vec<Array>) -> Result<Vec<Array>> {
24    Array::from_op(primitive, inputs)
25}
26
27/// The broadcast shape of `a` and `b` (NumPy rules), or `op`'s shape-mismatch error, as
28/// the binary operations compute it.
29pub fn broadcast_shapes(op: &'static str, a: &[usize], b: &[usize]) -> Result<Vec<usize>> {
30    Ok(crate::view::broadcast_shapes(op, a, b)?.to_vec())
31}
32
33/// A placeholder with this dtype and shape, standing for an input while a function is
34/// traced (as in MLX): operations can be recorded on it, but evaluating anything that
35/// depends on it fails. Placeholders are tracers.
36pub fn placeholder(dtype: DType, shape: &[i32]) -> Result<Array> {
37    let dims = shape
38        .iter()
39        .map(|&d| {
40            usize::try_from(d).map_err(|_| {
41                crate::Error::invalid_shape("placeholder", shape.iter().map(|&d| i64::from(d)))
42            })
43        })
44        .collect::<Result<Vec<_>>>()?;
45    Array::placeholder(dtype, &dims)
46}
47
48/// The arrays `a` is computed from, until it has been evaluated. Empty for host data and
49/// evaluated arrays.
50pub fn inputs(a: &Array) -> Vec<Array> {
51    match &*a.node.lock() {
52        State::Pending { inputs, .. } | State::Retained { inputs, .. } => inputs.clone(),
53        State::Running { .. } | State::Done | State::Failed { .. } => Vec::new(),
54    }
55}
56
57/// Every output of the operation that produces `a`, in order (`a` is the one at
58/// [`output_index`]).
59pub fn siblings(a: &Array) -> Vec<Array> {
60    (0..a.node.outputs.len() as u32)
61        .map(|output| Array {
62            node: a.node.clone(),
63            output,
64        })
65        .collect()
66}
67
68/// Which output of its operation `a` is.
69pub fn output_index(a: &Array) -> usize {
70    a.output as usize
71}
72
73/// Identifies the operation that produces `a` (shared by its siblings) for as long as the
74/// array is alive.
75pub fn id(a: &Array) -> usize {
76    std::sync::Arc::as_ptr(&a.node) as usize
77}
78
79/// Whether `a` depends on a transform's tracer or a placeholder, that is, it is being
80/// traced by `grad`, `vjp`, `jvp`, `vmap` or `compile`. Tracer status is transient: an
81/// array is a tracer only while the outermost transform it depends on runs (on this
82/// thread).
83pub fn is_tracer(a: &Array) -> bool {
84    a.node.is_tracer()
85}
86
87/// The operation that produces `a`, until it has been evaluated. `None` for host data and
88/// evaluated arrays.
89pub fn primitive(a: &Array) -> Option<Primitive> {
90    match &*a.node.lock() {
91        State::Pending { primitive, .. } => Some(primitive.clone()),
92        _ => None,
93    }
94}
95
96/// Per-dimension steps of `a`'s layout, in elements.
97pub fn strides(a: &Array) -> &[isize] {
98    &a.out().strides
99}
100
101/// Memory-layout flags of `a`.
102pub fn flags(a: &Array) -> Flags {
103    a.out().flags
104}
105
106/// `a[args]` with one [`SliceArg`] per dimension (Python semantics; a view), the lowered
107/// form of basic indexing.
108pub fn slice(a: &Array, args: &[SliceArg]) -> Result<Array> {
109    crate::ops::imp::slice(a, args)
110}
nuvec_core/graph.rs

nuvec_core/
graph.rs
