Crate constensor_core

Constensor is an experimental ML framework featuring a graph-based JIT compiler.

It’s designed with clarity and efficiency in mind, so besides compile-time checks for shape, dtype, and device, it also is based around a graph design.

This means that all the code you write is entered into a graph and then analyzed when you call Graph::optimize! The optimization step fuses operations, allows for automatic and seamless inplacing, constant folding, and other features.

Then, by precompliling the graph with Graph::compile, we can make intelligent decisions about when we run certain operations. For instance, on CUDA, streams are automatically used where possible to parallelize execution.

Currently, only CUDA and CPU are supported but Metal and cubecl is support coming very soon.

A quick guide

• First, create a Graph. This will hold all the operations and do optimization and compilation.
• Tensors are modelled with a GraphTensor. These represent the operation but do not perform any computation.
• Be sure to optimize the graph using Graph::optimize after all operations are complete!
• Compile the graph using Graph::compile, which will insert the device-specific optimizations. This returns a CompiledGraph.
• Run using CompiledGraph::run. This returns a concrete Tensor.

What can you do with it?

use constensor_core::{Cpu, Graph, GraphTensor, Tensor, R1, R2};

let mut graph: Graph<f32> = Graph::empty();
let _arange = GraphTensor::<R1<10>, f32, Cpu>::arange(&mut graph, 0., 1.);
let a = GraphTensor::<R2<3, 4>, f32, Cpu>::fill(&mut graph, 1.0);
let b = GraphTensor::<R2<3, 4>, f32, Cpu>::fill(&mut graph, 2.0);
let c = GraphTensor::<R2<3, 4>, f32, Cpu>::fill(&mut graph, 3.0);
let d = GraphTensor::<R2<3, 4>, f32, Cpu>::fill(&mut graph, 4.0);
let res = a * b + c;
let _out = res + d;

graph.optimize();

let compiled: constensor_core::CompiledGraph<R2<3, 4>, f32, Cpu> = graph.compile().unwrap();
let res = compiled.run().unwrap();

let tensor: Tensor<R2<3, 4>, f32, Cpu> = res;

assert_eq!(tensor.data().unwrap().to_vec(), vec![vec![9.0; 4]; 3],);

Macros

bail

Structs

Cpu

Graph

GraphNode

GraphTensor

A tensor representing an intermediary result of a graph. Performing operations on this tensor will not cause any computations.

R1

R2

R3

R4

R5

R6

Tensor

Tensors are n dimensional arrays. Only functions which allocate, copy data, or do operations return Results.

Enums

CompiledGraph

A representation of the compiled graph. The shape is the output shape.

Error

Op

Traits

Context

Attach more context to an error.

DType

Marker trait for tensor datatypes.

Shape

Marker trait for shapes

Type Aliases

BestDevice

Result