1pub mod clamped_swiglu;
2mod data_formats;
3pub mod gelu_approximate;
4pub mod gelu_exact;
5pub mod grid_sample;
6mod reduce;
7pub mod resize;
8pub mod rms_norm;
9pub mod silu;
10mod softmax;
11
12pub use self::clamped_swiglu::ClampedSwiGlu;
13pub use self::data_formats::{BaseDataShape, DataFormat, DataShape, SymDataShape};
14pub use self::gelu_approximate::GeluApproximate;
15pub use self::gelu_exact::GeluExact;
16pub use self::grid_sample::{GridSample, InterpolationMode, PaddingMode};
17pub use self::reduce::{Reduce, Reducer, expand_mean_of_squares};
18pub use self::resize::{
19 CoordTransformer, Interpolator, Nearest, NearestUpsample, Resize,
20 rewrite_nearest_upsample_to_broadcast,
21};
22pub use self::rms_norm::{RmsNorm, ScaledRmsNorm};
23pub use self::silu::Silu;
24pub use self::softmax::{Softmax, SoftmaxKind};
25
26pub use crate::internal::*;
27
28use tract_linalg::routines::Func;
29use tract_num_traits::AsPrimitive;
30
31element_wise!(sigmoid, Sigmoid,
32 [f16] => |_, xs| { Func::Sigmoid.ew_f16()?.run(xs) },
33 [f32] => |_, xs| { Func::Sigmoid.ew_f32()?.run(xs) };
34 q: [i8, u8, i32, i32] => |x: f32| 1.0 / (1.0+(-x).exp());
35 cost: |dt| {tvec!((Cost::FMA(dt), 11), (Cost::Div(dt), 1))};
36 declutter: silu::detect_silu
37);
38
39element_wise!(hard_swish, HardSwish,
40[f16] => |_, xs| { xs.iter_mut().for_each(|x| *x = *x * f16::from_f32(0.0).max(f16::from_f32(1.0).min(f16::from_f32(1. / 6.) * *x + f16::from_f32(0.5)))); Ok(()) },
41[f32] => |_, xs| { Func::Hardswish.ew_f32()?.run(xs) }
42 );
43
44element_wise!(leaky_relu, LeakyRelu { alpha: f32 },
45 [f16] => |op, xs| { Func::LeakyRelu.ew_f16_param()?.run_with_params(xs, f16::from_f32(op.alpha)) },
46 [f32] => |op, xs| { Func::LeakyRelu.ew_f32_param()?.run_with_params(xs, op.alpha) }
47);