Merge pull request jax-ml#2642 from jakevdp/trunc

Add implementation of np.trunc

Author: mattjj

jax/numpy/lax_numpy.py

@@ -653,6 +653,12 @@ def signbit(x):
   x = lax.bitcast_convert_type(x, int_type)
   return lax.convert_element_type(x >> (info.nexp + info.nmant), onp.bool)
 
+
+@_wraps(onp.trunc)
+def trunc(x):
+  return where(lax.lt(x, lax._const(x, 0)), lax.ceil(x), lax.floor(x))
+
+
 def _normalize_float(x):
     info = finfo(_dtype(x))
     cond = lax.abs(x) < info.tiny

tests/lax_numpy_test.py

@@ -132,6 +132,7 @@ def op_record(name, nargs, dtypes, shapes, rng_factory, diff_modes,
     op_record("subtract", 2, number_dtypes, all_shapes, jtu.rand_default, ["rev"]),
     op_record("signbit", 1, default_dtypes + bool_dtypes, all_shapes,
               jtu.rand_some_inf_and_nan, ["rev"]),
+    op_record("trunc", 1, float_dtypes, all_shapes, jtu.rand_some_inf_and_nan, []),
     op_record("sin", 1, number_dtypes, all_shapes, jtu.rand_default, ["rev"],
               inexact=True),
     op_record("cos", 1, number_dtypes, all_shapes, jtu.rand_default, ["rev"],