Merge branch 'main' of https://github.com/jafioti/luminal

crates/luminal_cuda/src/binary.rs

@@ -85,6 +85,7 @@ impl<T: CudaFloat> Operator for CudaSub<T> {
 pub struct SubtractionCompiler<T: CudaFloat>(PhantomData<T>);
 
 impl<T: CudaFloat> Compiler for SubtractionCompiler<T> {
+    type Output = ();
     fn compile<To: ToIdsMut>(&self, graph: &mut Graph, _: To) {
         let dev = CudaDevice::new(0).unwrap();
         let (lhs, rhs) = (node(), node());
@@ -208,6 +209,7 @@ impl<T: CudaFloat> Operator for CudaEqual<T> {
 pub struct EqualCompiler<T: CudaFloat>(PhantomData<T>);
 
 impl<T: CudaFloat> Compiler for EqualCompiler<T> {
+    type Output = ();
     fn compile<To: ToIdsMut>(&self, graph: &mut Graph, _: To) {
         let dev = CudaDevice::new(0).unwrap();
         let one = constant::<T>(1.);
@@ -343,6 +345,7 @@ impl<T: CudaFloat> Operator for CudaGather<T> {
 pub struct GatherCompiler<T: CudaFloat>(PhantomData<T>);
 
 impl<T: CudaFloat> Compiler for GatherCompiler<T> {
+    type Output = ();
     fn compile<To: ToIdsMut>(&self, graph: &mut Graph, _: To) {
         let dev = CudaDevice::new(0).unwrap();
         let indexes = node();

crates/luminal_cuda/src/matmul.rs

@@ -109,6 +109,7 @@ impl<T: CudaFloat + 'static> Compiler for MatMulCompiler<T>
 where
     CudaData<T>: Data,
 {
+    type Output = ();
     fn compile<To: ToIdsMut>(&self, graph: &mut Graph, mut ids: To) {
         let dev = CudaDevice::new(0).unwrap();
         // Look for the matmul pattern

crates/luminal_cuda/src/other.rs

@@ -75,6 +75,7 @@ impl<T: CudaFloat> Operator for CudaARange<T> {
 pub struct ARangeCompiler<T: CudaFloat>(PhantomData<T>);
 
 impl<T: CudaFloat> Compiler for ARangeCompiler<T> {
+    type Output = ();
     fn compile<To: ToIdsMut>(&self, graph: &mut Graph, _: To) {
         let dev = CudaDevice::new(0).unwrap();
         // TODO: Make sure this actually checks the shape transformations to ensure pooling happens

crates/luminal_cuda/src/prim.rs

@@ -937,6 +937,7 @@ impl<T: CudaFloat> Operator for CudaMaxReduce<T> {
 pub struct PrimitiveCompiler<T>(PhantomData<T>);
 
 impl<T: CudaFloat> Compiler for PrimitiveCompiler<T> {
+    type Output = ();
     fn compile<To: ToIdsMut>(&self, graph: &mut Graph, mut ids: To) {
         let dev = CudaDevice::new(0).unwrap();
         // Go through the graph and insert copy ops
@@ -1146,6 +1147,7 @@ impl<T: CudaFloat> Compiler for PrimitiveCompiler<T> {
 pub struct CopyCompiler<T>(PhantomData<T>);
 
 impl<T: CudaFloat> Compiler for CopyCompiler<T> {
+    type Output = ();
     fn compile<To: ToIdsMut>(&self, graph: &mut Graph, mut ids: To) {
         for (first, second) in graph
             .edge_indices()

crates/luminal_cuda/src/quantized.rs

@@ -259,6 +259,7 @@ impl<T> CudaQuantizedCompiler<T> {
 }
 
 impl<T: CudaFloat + Default> Compiler for CudaQuantizedCompiler<T> {
+    type Output = ();
     fn compile<To: ToIdsMut>(&self, graph: &mut Graph, mut remap: To) {
         let device = CudaDevice::new(0).unwrap();
         let mut weight_ids = self.0.clone();

examples/llama/src/model.rs

@@ -11,7 +11,6 @@ pub const HEADS: usize = 32;
 pub const LAYERS: usize = 32;
 
 use luminal::{
-    nn::{embedding::Embedding, norm::RMSNorm},
     prelude::*,
     shape::symbolic::{BigExpression, Expression},
 };

examples/mistral/src/model.rs

@@ -1,7 +1,6 @@
 use std::{marker::PhantomData, ops::Div};
 
 use luminal::{
-    nn::{embedding::Embedding, norm::RMSNorm},
     prelude::{binary::F32Pow, *},
     shape::symbolic::{BigExpression, Expression},
 };

examples/simple/src/main.rs

@@ -1,4 +1,4 @@
-use luminal::{nn::linear::Linear, prelude::*};
+use luminal::prelude::*;
 
 fn main() {
     // Create a new graph