Add Exp as a core op

catgrad/src/category/core.rs

@@ -256,6 +256,7 @@ pub enum ScalarOp {
     EQ,    // 2 → 1
     Cos,   // 1 → 1
     Sin,   // 1 → 1
+    Exp,   // 1 → 1
     Log,   // 1 → 1
     Floor, // 1 → 1
     Where, // 3 → 1
@@ -278,6 +279,7 @@ impl ScalarOp {
             ScalarOp::EQ => (2, 1),
             ScalarOp::Cos => (1, 1),
             ScalarOp::Sin => (1, 1),
+            ScalarOp::Exp => (1, 1),
             ScalarOp::Log => (1, 1),
             ScalarOp::Floor => (1, 1),
             ScalarOp::Where => (3, 1),

catgrad/src/category/lang/ops.rs

@@ -88,6 +88,10 @@ pub fn cos(builder: &Builder, value: Var) -> Var {
     var::fn_operation(builder, &[value], Object::Tensor, op!["tensor", "cos"])
 }
 
+pub fn exp(builder: &Builder, value: Var) -> Var {
+    var::fn_operation(builder, &[value], Object::Tensor, op!["tensor", "exp"])
+}
+
 pub fn log(builder: &Builder, value: Var) -> Var {
     var::fn_operation(builder, &[value], Object::Tensor, op!["tensor", "log"])
 }

catgrad/src/interpreter/backend/candle.rs

@@ -792,6 +792,17 @@ impl Backend for CandleBackend {
         }
     }
 
+    fn exp(&self, x: TaggedTensor<Self>) -> TaggedTensor<Self> {
+        use TaggedTensorTuple::*;
+        match x {
+            F32([arr]) => F32([Self::unary_eager(arr, DType::F32, Self::exp)]),
+            F16([arr]) => F16([Self::unary_eager(arr, DType::F16, Self::exp)]),
+            BF16([arr]) => BF16([Self::unary_eager(arr, DType::BF16, Self::exp)]),
+            FP8([arr]) => FP8([Self::unary_eager(arr, DType::F8E4M3, Self::exp)]),
+            _ => panic!("Invalid type for exp"),
+        }
+    }
+
     fn log(&self, x: TaggedTensor<Self>) -> TaggedTensor<Self> {
         use TaggedTensorTuple::*;
         match x {
@@ -1341,6 +1352,10 @@ impl CandleBackend {
         x.cos().unwrap().into()
     }
 
+    fn exp(x: &Tensor) -> CandleTensor {
+        x.exp().unwrap().into()
+    }
+
     fn log(x: &Tensor) -> CandleTensor {
         x.log().unwrap().into()
     }

catgrad/src/interpreter/backend/mod.rs

@@ -72,6 +72,7 @@ pub trait Backend: Clone + Debug {
     fn where_cond(&self, args: TaggedTensorTuple<Self, 3>) -> TaggedTensor<Self>;
     fn sin(&self, x: TaggedTensor<Self>) -> TaggedTensor<Self>;
     fn cos(&self, x: TaggedTensor<Self>) -> TaggedTensor<Self>;
+    fn exp(&self, x: TaggedTensor<Self>) -> TaggedTensor<Self>;
     fn log(&self, x: TaggedTensor<Self>) -> TaggedTensor<Self>;
     fn floor(&self, x: TaggedTensor<Self>) -> TaggedTensor<Self>;
     fn neg(&self, x: TaggedTensor<Self>) -> TaggedTensor<Self>;

catgrad/src/interpreter/backend/ndarray.rs

@@ -374,6 +374,14 @@ impl Backend for NdArrayBackend {
         }
     }
 
+    fn exp(&self, x: TaggedTensor<Self>) -> TaggedTensor<Self> {
+        use TaggedTensorTuple::*;
+        match x {
+            F32([arr]) => from_f32(arr.unwrap_f32().exp()),
+            _ => panic!("Invalid input types for exp"),
+        }
+    }
+
     fn log(&self, x: TaggedTensor<Self>) -> TaggedTensor<Self> {
         use TaggedTensorTuple::*;
         match x {

catgrad/src/interpreter/backend/shape_only.rs

@@ -189,6 +189,10 @@ impl Backend for ShapeOnlyBackend {
         x
     }
 
+    fn exp(&self, x: TaggedTensor<Self>) -> TaggedTensor<Self> {
+        x
+    }
+
     fn log(&self, x: TaggedTensor<Self>) -> TaggedTensor<Self> {
         x
     }

catgrad/src/interpreter/tensor_op.rs

@@ -18,6 +18,7 @@ pub(crate) fn tensor_op<B: Backend>(
         TensorOp::Map(ScalarOp::Pow) => binop(backend, args, ssa, B::pow),
         TensorOp::Map(ScalarOp::Sin) => unary_op(backend, args, ssa, B::sin),
         TensorOp::Map(ScalarOp::Cos) => unary_op(backend, args, ssa, B::cos),
+        TensorOp::Map(ScalarOp::Exp) => unary_op(backend, args, ssa, B::exp),
         TensorOp::Map(ScalarOp::Log) => unary_op(backend, args, ssa, B::log),
         TensorOp::Map(ScalarOp::Floor) => unary_op(backend, args, ssa, B::floor),
         TensorOp::Map(ScalarOp::Neg) => unary_op(backend, args, ssa, B::neg),

catgrad/src/pass/to_core.rs

@@ -102,6 +102,7 @@ pub(crate) fn core_declarations() -> HashMap<lang::Path, core::Operation> {
         (path!["tensor", "pow"], Operation::Tensor(Map(Pow))),
         (path!["tensor", "sin"], Operation::Tensor(Map(Sin))),
         (path!["tensor", "cos"], Operation::Tensor(Map(Cos))),
+        (path!["tensor", "exp"], Operation::Tensor(Map(Exp))),
         (path!["tensor", "log"], Operation::Tensor(Map(Log))),
         (path!["tensor", "floor"], Operation::Tensor(Map(Floor))),
         (path!["tensor", "lt"], Operation::Tensor(Map(LT))),

catgrad/src/stdlib/nn/mod.rs

@@ -5,7 +5,7 @@
 //! higher level LLM building blocks or APIs.
 
 use crate::prelude::{ops::*, *};
-use std::f32::consts::{E, PI};
+use std::f32::consts::PI;
 
 use crate::typecheck::*;
 
@@ -121,13 +121,6 @@ pub fn sqrt(builder: &Builder, x: Var) -> Var {
     pow(builder, x, e)
 }
 
-pub fn exp(builder: &Builder, x: Var) -> Var {
-    let sh = shape(builder, x.clone());
-    let e = constant(builder, E, &sh);
-    let e = cast(builder, e, dtype(builder, x.clone()));
-    pow(builder, e, x)
-}
-
 pub fn sigmoid(builder: &Builder, x: Var) -> Var {
     let sh = shape(builder, x.clone());
     let one = constant(builder, 1.0, &sh);

catgrad/src/stdlib/ops.rs

@@ -4,9 +4,9 @@ use crate::prelude::{Builder, Var};
 
 // re-export lang ops
 pub use ops::{
-    arange, argmax, broadcast, cast, concat, cond, cos, dtype, dtype_constant, eq, floor, gt, gte,
-    index, log, lt, lte, matmul, max, nat, nat_to_u32, pack, param, pow, probe, reshape, shape,
-    sin, slice, sum, topk, transpose, unpack, where_cond,
+    arange, argmax, broadcast, cast, concat, cond, cos, dtype, dtype_constant, eq, exp, floor, gt,
+    gte, index, log, lt, lte, matmul, max, nat, nat_to_u32, pack, param, pow, probe, reshape,
+    shape, sin, slice, sum, topk, transpose, unpack, where_cond,
 };
 
 pub fn get(builder: &Builder, dim: impl IntoNatVar, start: impl IntoNatVar, x: Var) -> Var {