Integrate Triton up to [bfb8e413](https://github.com/openai/triton/co…

…mmits/bfb8e413b075583228c961bdbb65a98dc54d0868) PiperOrigin-RevId: 612852008
jax-ml · Mar 7, 2024 · ed94232 · ed94232
1 parent 08f1063
commit ed94232
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/tests/triton_test.py b/tests/triton_test.py
@@ -20,6 +20,7 @@
 import numpy as np
 import triton
 import triton.language as tl
+from triton.language.extra.cuda import libdevice
 
 
 @triton.jit
@@ -71,7 +72,7 @@ def tanh_kernel(
   # Load x and y from DRAM, masking out any extra elements in case the input is not a
   # multiple of the block size
   x = tl.load(x_ptr + offsets, mask=mask)
-  output = tl.math.tanh(x)
+  output = libdevice.tanh(x)
   # Write x + y back to DRAM
   tl.store(output_ptr + offsets, output, mask=mask)