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/// A placeholder with this dtype and shape, standing for an input while a function is
28/// traced (as in MLX): operations can be recorded on it, but evaluating anything that
29/// depends on it fails. Placeholders are tracers.
30pub fn placeholder(dtype: DType, shape: &[i32]) -> Result<Array> {
31    let dims = shape
32        .iter()
33        .map(|&d| {
34            usize::try_from(d).map_err(|_| {
35                crate::Error::invalid_shape("placeholder", shape.iter().map(|&d| i64::from(d)))
36            })
37        })
38        .collect::<Result<Vec<_>>>()?;
39    Array::placeholder(dtype, &dims)
40}
41
42/// The arrays `a` is computed from, until it has been evaluated. Empty for host data and
43/// evaluated arrays.
44pub fn inputs(a: &Array) -> Vec<Array> {
45    match &*a.node.lock() {
46        State::Pending { inputs, .. } | State::Retained { inputs, .. } => inputs.clone(),
47        State::Running { .. } | State::Done | State::Failed { .. } => Vec::new(),
48    }
49}
50
51/// Every output of the operation that produces `a`, in order (`a` is the one at
52/// [`output_index`]).
53pub fn siblings(a: &Array) -> Vec<Array> {
54    (0..a.node.outputs.len() as u32)
55        .map(|output| Array {
56            node: a.node.clone(),
57            output,
58        })
59        .collect()
60}
61
62/// Which output of its operation `a` is.
63pub fn output_index(a: &Array) -> usize {
64    a.output as usize
65}
66
67/// Identifies the operation that produces `a` (shared by its siblings) for as long as the
68/// array is alive.
69pub fn id(a: &Array) -> usize {
70    std::sync::Arc::as_ptr(&a.node) as usize
71}
72
73/// Whether `a` depends on a transform's tracer or a placeholder, that is, it is being
74/// traced by `grad`, `vjp`, `jvp`, `vmap` or `compile`. Tracer status is transient: an
75/// array is a tracer only while the outermost transform it depends on runs (on this
76/// thread).
77pub fn is_tracer(a: &Array) -> bool {
78    a.node.is_tracer()
79}
80
81/// The operation that produces `a`, until it has been evaluated. `None` for host data and
82/// evaluated arrays.
83pub fn primitive(a: &Array) -> Option<Primitive> {
84    match &*a.node.lock() {
85        State::Pending { primitive, .. } => Some(primitive.clone()),
86        _ => None,
87    }
88}
89
90/// Per-dimension steps of `a`'s layout, in elements.
91pub fn strides(a: &Array) -> &[isize] {
92    &a.out().strides
93}
94
95/// Memory-layout flags of `a`.
96pub fn flags(a: &Array) -> Flags {
97    a.out().flags
98}
99
100/// `a[args]` with one [`SliceArg`] per dimension (Python semantics; a view), the lowered
101/// form of basic indexing.
102pub fn slice(a: &Array, args: &[SliceArg]) -> Result<Array> {
103    crate::ops::imp::slice(a, args)
104}
nuvec_core/graph.rs

nuvec_core/
graph.rs
