dfm
commented
3 weeks ago I don't know the answer to this, but maybe @jakevdp does?
Some notes in the meantime. It's worth checking out the JAX microbenchmark FAQ entry because benchmarking like you're doing here can lead to incorrect conclusions since it includes the tracing and compilation overhead. Updating this doesn't seem to change the specific conclusions though! Here's how I would write the benchmark:
```python
import numpy as np
import jax
import jax.numpy as jnp
num_segments = 1700
segment_ids = np.repeat(np.arange(num_segments), np.random.randint(40,977,size=num_segments))
@jax.jit
def do_sum(data):
return jax.ops.segment_sum(data, segment_ids=segment_ids, num_segments=num_segments, indices_are_sorted=True)
key = jax.random.PRNGKey(0)
data = jax.random.uniform(key, shape=(len(segment_ids),), dtype=jnp.float32)
do_sum(data).block_until_ready() # compile
%timeit do_sum(data).block_until_ready()
key = jax.random.PRNGKey(0)
data = jax.random.uniform(key, shape=(len(segment_ids),), dtype=jnp.float16)
do_sum(data).block_until_ready() # compile
%timeit do_sum(data).block_until_ready()
```
Regardless, I do find that the float16 version is consistently slower. Perhaps @jakevdp can lead us in the right direction!
jakevdp
CloudyDory
Description
I find that JAX segment_sum is two times slower for FP16 inputs than FP32 inputs. Here is an example:
Outputs:
This happens with or without
jit(). Why does this happen? And is there a way to optimize the computation for FP16 input?System info (python version, jaxlib version, accelerator, etc.)