Merge pull request #49 from BradenEverson/bf16

is_zero checks inner constant data type not node data type
unda-ml · Mar 9, 2024 · 9cea4df · 9cea4df
2 parents 5ff81a2 + 99880c8
commit 9cea4df
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 4 deletions.
diff --git a/src/core/graph/compile.rs b/src/core/graph/compile.rs
@@ -32,9 +32,11 @@ impl Context {
             Err(CompileError::NoReturn)?;
         }
 
+
         for a in returns.iter() {
             self.fold_consts(*a, usize::MAX)?;
         }
+
         //while self.foldconsts(a, 1)? {
         //    println!("{}", self.to_string(a));
         //}
@@ -419,7 +421,6 @@ impl Context {
                 }
             }
         }
-
         let xla_return_vec: Vec<&xla::XlaOp> = returns
             .into_iter()
             .map(|i| &xla_op_slotmap[unda_xla_map[&i.into()]])

diff --git a/src/core/graph/mod.rs b/src/core/graph/mod.rs
@@ -11,7 +11,7 @@ mod operation;
 mod parameter;
 mod shape;
 mod subterm;
-mod tests;
+//mod tests;
 mod tests_cpu;
 
 use callsite::{callsite, Callsite};

diff --git a/src/core/graph/node.rs b/src/core/graph/node.rs
@@ -7,7 +7,7 @@ use xla::Literal;
 use std::{fmt::{Display, Formatter, Result}, error::Error};
 
 /// A node in the compute graph
-#[derive(Clone)]
+#[derive(Clone, Debug)]
 pub struct Node {
     /// helps identify where in the user's source code this node originated
     // TODO: gate this so its not present at all in release builds
@@ -39,7 +39,7 @@ impl Node {
         //TODO! Convert type to primative type so we can collect the values
         return match &self.operation {
             Operation::Constant(a) => {
-                match self.dtype {
+                match a.value.element_type()? {
                     xla::ElementType::F32 => {
                         let data_ref = a.value.to_vec::<f32>()?;
                         for i in data_ref.iter() {