FEATURE: clamp function in arith module

Author: 9prady9

src/arith/mod.rs

@@ -42,6 +42,7 @@ extern {
     fn af_bitshiftr(out: MutAfArray, lhs: AfArray, rhs: AfArray, batch: c_int) -> c_int;
     fn af_minof(out: MutAfArray, lhs: AfArray, rhs: AfArray, batch: c_int) -> c_int;
     fn af_maxof(out: MutAfArray, lhs: AfArray, rhs: AfArray, batch: c_int) -> c_int;
+    fn af_clamp(out: MutAfArray, inp: AfArray, lo: AfArray, hi: AfArray, batch: c_int) -> c_int;
 
     fn af_not(out: MutAfArray, arr: AfArray) -> c_int;
     fn af_abs(out: MutAfArray, arr: AfArray) -> c_int;
@@ -293,6 +294,35 @@ overloaded_binary_func!("Create complex array from two Arrays", cplx2, cplx2_hel
 overloaded_binary_func!("Compute root", root, root_helper, af_root);
 overloaded_binary_func!("Computer power", pow, pow_helper, af_pow);
 
+pub fn clamp<T, U> (input: &Array, arg1: &T, arg2: &U, batch: bool) -> Array
+    where T: Convertable, U: Convertable
+{
+    let clamp_helper = |lo: &Array, hi: &Array| {
+        unsafe {
+            let mut temp: i64 = 0;
+            let err_val = af_clamp(&mut temp as MutAfArray, input.get() as AfArray,
+                                   lo.get() as AfArray, hi.get() as AfArray,
+                                   batch as c_int);
+            HANDLE_ERROR(AfError::from(err_val));
+            Array::from(temp)
+        }
+    };
+
+    let lo = arg1.convert();
+    let hi = arg2.convert();
+    match (lo.is_scalar(), hi.is_scalar()) {
+        ( true, false) => {
+            let l = tile(&lo, hi.dims());
+            clamp_helper(&l, &hi)
+        },
+        (false,  true) => {
+            let r = tile(&hi, lo.dims());
+            clamp_helper(&lo, &r)
+        },
+        _ => clamp_helper(&lo, &hi),
+    }
+}
+
 macro_rules! arith_scalar_func {
     ($rust_type: ty, $op_name:ident, $fn_name: ident, $ffi_fn: ident) => (
         impl<'f> $op_name<$rust_type> for &'f Array {

src/lib.rs

@@ -19,7 +19,7 @@ mod algorithm;
 
 pub use arith::{add, sub, div, mul, lt, gt, le, ge, eq, neq, and, or, minof, maxof, rem};
 pub use arith::{bitand, bitor, bitxor, shiftl, shiftr};
-pub use arith::{abs, sign, round, trunc, floor, ceil, modulo, sigmoid};
+pub use arith::{abs, sign, round, trunc, floor, ceil, modulo, sigmoid, clamp};
 pub use arith::{sin, cos, tan, asin, acos, atan, sinh, cosh, tanh, asinh, acosh, atanh};
 pub use arith::{atan2, cplx2, arg, cplx, real, imag, conjg, hypot};
 pub use arith::{sqrt, log, log1p, log10, log2, pow2, exp, expm1, erf, erfc, root, pow};