diff --git a/catgrad/src/category/core.rs b/catgrad/src/category/core.rs index 62c81ea8..ce0bf2f6 100644 --- a/catgrad/src/category/core.rs +++ b/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), diff --git a/catgrad/src/category/lang/ops.rs b/catgrad/src/category/lang/ops.rs index c774a479..ac41c7fe 100644 --- a/catgrad/src/category/lang/ops.rs +++ b/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"]) } diff --git a/catgrad/src/interpreter/backend/candle.rs b/catgrad/src/interpreter/backend/candle.rs index a12972b6..60c7ce5a 100644 --- a/catgrad/src/interpreter/backend/candle.rs +++ b/catgrad/src/interpreter/backend/candle.rs @@ -792,6 +792,17 @@ impl Backend for CandleBackend { } } + fn exp(&self, x: TaggedTensor) -> TaggedTensor { + 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) -> TaggedTensor { 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() } diff --git a/catgrad/src/interpreter/backend/mod.rs b/catgrad/src/interpreter/backend/mod.rs index b7ae6721..2e7e69b4 100644 --- a/catgrad/src/interpreter/backend/mod.rs +++ b/catgrad/src/interpreter/backend/mod.rs @@ -72,6 +72,7 @@ pub trait Backend: Clone + Debug { fn where_cond(&self, args: TaggedTensorTuple) -> TaggedTensor; fn sin(&self, x: TaggedTensor) -> TaggedTensor; fn cos(&self, x: TaggedTensor) -> TaggedTensor; + fn exp(&self, x: TaggedTensor) -> TaggedTensor; fn log(&self, x: TaggedTensor) -> TaggedTensor; fn floor(&self, x: TaggedTensor) -> TaggedTensor; fn neg(&self, x: TaggedTensor) -> TaggedTensor; diff --git a/catgrad/src/interpreter/backend/ndarray.rs b/catgrad/src/interpreter/backend/ndarray.rs index fe0a4d74..e83a06ca 100644 --- a/catgrad/src/interpreter/backend/ndarray.rs +++ b/catgrad/src/interpreter/backend/ndarray.rs @@ -374,6 +374,14 @@ impl Backend for NdArrayBackend { } } + fn exp(&self, x: TaggedTensor) -> TaggedTensor { + use TaggedTensorTuple::*; + match x { + F32([arr]) => from_f32(arr.unwrap_f32().exp()), + _ => panic!("Invalid input types for exp"), + } + } + fn log(&self, x: TaggedTensor) -> TaggedTensor { use TaggedTensorTuple::*; match x { diff --git a/catgrad/src/interpreter/backend/shape_only.rs b/catgrad/src/interpreter/backend/shape_only.rs index cab4f3a3..97eb6e8a 100644 --- a/catgrad/src/interpreter/backend/shape_only.rs +++ b/catgrad/src/interpreter/backend/shape_only.rs @@ -189,6 +189,10 @@ impl Backend for ShapeOnlyBackend { x } + fn exp(&self, x: TaggedTensor) -> TaggedTensor { + x + } + fn log(&self, x: TaggedTensor) -> TaggedTensor { x } diff --git a/catgrad/src/interpreter/tensor_op.rs b/catgrad/src/interpreter/tensor_op.rs index 9e353154..97b06d0a 100644 --- a/catgrad/src/interpreter/tensor_op.rs +++ b/catgrad/src/interpreter/tensor_op.rs @@ -18,6 +18,7 @@ pub(crate) fn tensor_op( 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), diff --git a/catgrad/src/pass/to_core.rs b/catgrad/src/pass/to_core.rs index 77d86231..5a4e5136 100644 --- a/catgrad/src/pass/to_core.rs +++ b/catgrad/src/pass/to_core.rs @@ -102,6 +102,7 @@ pub(crate) fn core_declarations() -> HashMap { (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))), diff --git a/catgrad/src/stdlib/nn/mod.rs b/catgrad/src/stdlib/nn/mod.rs index ca17f9ce..82800e3b 100644 --- a/catgrad/src/stdlib/nn/mod.rs +++ b/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); diff --git a/catgrad/src/stdlib/ops.rs b/catgrad/src/stdlib/ops.rs index e8bb2cd6..415b6f7b 100644 --- a/catgrad/src/stdlib/ops.rs +++ b/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 {