Understanding the source code of jax.lax | tanh

[Rohanjames1997]

Hi there,

I'm trying to understand how the tanh function is implemented in jax, and I'm interested in knowing where the actual underlying "math" resides in the codebase.

Does it simply call numpy's/XLA's tanh operation? If I were to optimize the tanh op in jax using SIMD intrinsics, where should I be looking to do so?

So far, from the repo, I've gotten this, but I couldn't make much of it :

jax/jax/_src/lax/lax.py

Lines 337 to 339 in 8ef41a6

	def tanh(x: ArrayLike) -> Array:
	r"""Elementwise hyperbolic tangent: :math:`\mathrm{tanh}(x)`."""
	return tanh_p.bind(x)

Any pointers appreciated, TIA!

[dfm]

It's not so straightforward to follow the full stack, but the basic flow is that a JAX function get's lowered to StableHLO, which you can inspect as follows:

import jax
import jax.numpy as jnp

print(jax.jit(jnp.tanh).lower(jnp.linspace(-1, 1, 5)).as_text())

Since StableHLO has a tanh intrinsic, the intermediate representation is (relatively) simple:

module @jit_tanh attributes {mhlo.num_partitions = 1 : i32, mhlo.num_replicas = 1 : i32} {
  func.func public @main(%arg0: tensor<5xf32> {mhlo.layout_mode = "default"}) -> (tensor<5xf32> {jax.result_info = "", mhlo.layout_mode = "default"}) {
    %0 = stablehlo.tanh %arg0 : tensor<5xf32>
    return %0 : tensor<5xf32>
  }
}

Then, depending on the specific platform that you are targeting, XLA will compile this program differently. For example, on CPU, I believe the tanh will be passed to LLVM, which (I expect) will emit SIMD operations.

So, the tl;dr is that it's not really straightforward to track down exactly how specific operations are compiled. If you want to contribute optimizations like this, those would live in LLVM or XLA rather than JAX itself.

Hope this helps!

[Rohanjames1997]

Thank you - this helps a lot!

Indeed - while debugging I found that on a CPU, the default XLA implementation of float32 tanh on aarch64 is in EmitFastTanh which in-turn emits LLVM SIMD ops like you said.