<p dir="auto">I have observed that there is large increase in computation time when differentiating through a loss that involves index a complex128-valued array. Consider the following demo script.</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="import argparse
import time

import jax.numpy as jnp
from jax import jit, grad
from jax import random
from jax.config import config


parser = argparse.ArgumentParser(description='Complex indexing and gradients in JAX')
parser.add_argument('--complex', default=True, dest='complex', action='store_true', help='Index a complex-valued object')
parser.add_argument('--no-complex', dest='complex', action='store_false')
parser.add_argument('--double', default=True, dest='double', action='store_true', help='Use double precision')
parser.add_argument('--no-double', dest='double', action='store_false')
parser.add_argument('--index', default=True, dest='index', action='store_true', help='Index the object inside the loss')
parser.add_argument('--no-index', dest='index', action='store_false')
args = parser.parse_args()

if args.double:
    config.update(&quot;jax_enable_x64&quot;, True)

def loss(fvol, idx):
    # Define a loss function that either indexes or does not index the array.
    if args.index:
        vv = fvol[idx].sum([-1])
        # Doing this is fast again in complex-128.
        # fvol_real = fvol.real[idx]
        # fvol_imag = fvol.imag[idx]
        # vv = jnp.sum(fvol_real + 1j*fvol_imag, [-1])
    else:
        vv = fvol
    return jnp.square(jnp.linalg.norm(vv))

@jit
def loss_grad(fvol, idx):
    return grad(loss)(fvol, idx)

# Create a large array.
rng = random.PRNGKey(0)
rng_real, rng_imag = random.split(rng, 2)
L = 256
fvol = random.normal(rng_real, [L**3])
if args.complex:
    fvol += 1j * random.normal(rng_imag, fvol.shape)

dtype = fvol.dtype
idx = random.randint(rng, [L*L*8], minval=0, maxval=L)

# Time how long it takes to compute the loss.
for i in range(3):
    start = time.time()
    lv = loss(fvol, idx)
    lv.block_until_ready()
    elapsed = time.time() - start
    print('object data: {} - time elapsed to evaluate loss: {:.5f}'.format(dtype, elapsed))

# Time how long it takes to compute the gradient of the loss.
for i in range(3):
    start = time.time()
    g = loss_grad(fvol, idx)
    g.block_until_ready()
    elapsed = time.time() - start
    print('object data: {} - time elapsed to evaluate gradient: {:.5f}'.format(dtype, elapsed))"><pre class="notranslate"><span class="pl-k">import</span> <span class="pl-s1">argparse</span>
<span class="pl-k">import</span> <span class="pl-s1">time</span>

<span class="pl-k">import</span> <span class="pl-s1">jax</span>.<span class="pl-s1">numpy</span> <span class="pl-k">as</span> <span class="pl-s1">jnp</span>
<span class="pl-k">from</span> <span class="pl-s1">jax</span> <span class="pl-k">import</span> <span class="pl-s1">jit</span>, <span class="pl-s1">grad</span>
<span class="pl-k">from</span> <span class="pl-s1">jax</span> <span class="pl-k">import</span> <span class="pl-s1">random</span>
<span class="pl-k">from</span> <span class="pl-s1">jax</span>.<span class="pl-s1">config</span> <span class="pl-k">import</span> <span class="pl-s1">config</span>


<span class="pl-s1">parser</span> <span class="pl-c1">=</span> <span class="pl-s1">argparse</span>.<span class="pl-v">ArgumentParser</span>(<span class="pl-s1">description</span><span class="pl-c1">=</span><span class="pl-s">'Complex indexing and gradients in JAX'</span>)
<span class="pl-s1">parser</span>.<span class="pl-en">add_argument</span>(<span class="pl-s">'--complex'</span>, <span class="pl-s1">default</span><span class="pl-c1">=</span><span class="pl-c1">True</span>, <span class="pl-s1">dest</span><span class="pl-c1">=</span><span class="pl-s">'complex'</span>, <span class="pl-s1">action</span><span class="pl-c1">=</span><span class="pl-s">'store_true'</span>, <span class="pl-s1">help</span><span class="pl-c1">=</span><span class="pl-s">'Index a complex-valued object'</span>)
<span class="pl-s1">parser</span>.<span class="pl-en">add_argument</span>(<span class="pl-s">'--no-complex'</span>, <span class="pl-s1">dest</span><span class="pl-c1">=</span><span class="pl-s">'complex'</span>, <span class="pl-s1">action</span><span class="pl-c1">=</span><span class="pl-s">'store_false'</span>)
<span class="pl-s1">parser</span>.<span class="pl-en">add_argument</span>(<span class="pl-s">'--double'</span>, <span class="pl-s1">default</span><span class="pl-c1">=</span><span class="pl-c1">True</span>, <span class="pl-s1">dest</span><span class="pl-c1">=</span><span class="pl-s">'double'</span>, <span class="pl-s1">action</span><span class="pl-c1">=</span><span class="pl-s">'store_true'</span>, <span class="pl-s1">help</span><span class="pl-c1">=</span><span class="pl-s">'Use double precision'</span>)
<span class="pl-s1">parser</span>.<span class="pl-en">add_argument</span>(<span class="pl-s">'--no-double'</span>, <span class="pl-s1">dest</span><span class="pl-c1">=</span><span class="pl-s">'double'</span>, <span class="pl-s1">action</span><span class="pl-c1">=</span><span class="pl-s">'store_false'</span>)
<span class="pl-s1">parser</span>.<span class="pl-en">add_argument</span>(<span class="pl-s">'--index'</span>, <span class="pl-s1">default</span><span class="pl-c1">=</span><span class="pl-c1">True</span>, <span class="pl-s1">dest</span><span class="pl-c1">=</span><span class="pl-s">'index'</span>, <span class="pl-s1">action</span><span class="pl-c1">=</span><span class="pl-s">'store_true'</span>, <span class="pl-s1">help</span><span class="pl-c1">=</span><span class="pl-s">'Index the object inside the loss'</span>)
<span class="pl-s1">parser</span>.<span class="pl-en">add_argument</span>(<span class="pl-s">'--no-index'</span>, <span class="pl-s1">dest</span><span class="pl-c1">=</span><span class="pl-s">'index'</span>, <span class="pl-s1">action</span><span class="pl-c1">=</span><span class="pl-s">'store_false'</span>)
<span class="pl-s1">args</span> <span class="pl-c1">=</span> <span class="pl-s1">parser</span>.<span class="pl-en">parse_args</span>()

<span class="pl-k">if</span> <span class="pl-s1">args</span>.<span class="pl-s1">double</span>:
    <span class="pl-s1">config</span>.<span class="pl-en">update</span>(<span class="pl-s">"jax_enable_x64"</span>, <span class="pl-c1">True</span>)

<span class="pl-k">def</span> <span class="pl-en">loss</span>(<span class="pl-s1">fvol</span>, <span class="pl-s1">idx</span>):
    <span class="pl-c"># Define a loss function that either indexes or does not index the array.</span>
    <span class="pl-k">if</span> <span class="pl-s1">args</span>.<span class="pl-s1">index</span>:
        <span class="pl-s1">vv</span> <span class="pl-c1">=</span> <span class="pl-s1">fvol</span>[<span class="pl-s1">idx</span>].<span class="pl-en">sum</span>([<span class="pl-c1">-</span><span class="pl-c1">1</span>])
        <span class="pl-c"># Doing this is fast again in complex-128.</span>
        <span class="pl-c"># fvol_real = fvol.real[idx]</span>
        <span class="pl-c"># fvol_imag = fvol.imag[idx]</span>
        <span class="pl-c"># vv = jnp.sum(fvol_real + 1j*fvol_imag, [-1])</span>
    <span class="pl-k">else</span>:
        <span class="pl-s1">vv</span> <span class="pl-c1">=</span> <span class="pl-s1">fvol</span>
    <span class="pl-k">return</span> <span class="pl-s1">jnp</span>.<span class="pl-en">square</span>(<span class="pl-s1">jnp</span>.<span class="pl-s1">linalg</span>.<span class="pl-en">norm</span>(<span class="pl-s1">vv</span>))

<span class="pl-en">@<span class="pl-s1">jit</span></span>
<span class="pl-k">def</span> <span class="pl-en">loss_grad</span>(<span class="pl-s1">fvol</span>, <span class="pl-s1">idx</span>):
    <span class="pl-k">return</span> <span class="pl-en">grad</span>(<span class="pl-s1">loss</span>)(<span class="pl-s1">fvol</span>, <span class="pl-s1">idx</span>)

<span class="pl-c"># Create a large array.</span>
<span class="pl-s1">rng</span> <span class="pl-c1">=</span> <span class="pl-s1">random</span>.<span class="pl-v">PRNGKey</span>(<span class="pl-c1">0</span>)
<span class="pl-s1">rng_real</span>, <span class="pl-s1">rng_imag</span> <span class="pl-c1">=</span> <span class="pl-s1">random</span>.<span class="pl-en">split</span>(<span class="pl-s1">rng</span>, <span class="pl-c1">2</span>)
<span class="pl-v">L</span> <span class="pl-c1">=</span> <span class="pl-c1">256</span>
<span class="pl-s1">fvol</span> <span class="pl-c1">=</span> <span class="pl-s1">random</span>.<span class="pl-en">normal</span>(<span class="pl-s1">rng_real</span>, [<span class="pl-v">L</span><span class="pl-c1">**</span><span class="pl-c1">3</span>])
<span class="pl-k">if</span> <span class="pl-s1">args</span>.<span class="pl-s1">complex</span>:
    <span class="pl-s1">fvol</span> <span class="pl-c1">+=</span> <span class="pl-c1">1j</span> <span class="pl-c1">*</span> <span class="pl-s1">random</span>.<span class="pl-en">normal</span>(<span class="pl-s1">rng_imag</span>, <span class="pl-s1">fvol</span>.<span class="pl-s1">shape</span>)

<span class="pl-s1">dtype</span> <span class="pl-c1">=</span> <span class="pl-s1">fvol</span>.<span class="pl-s1">dtype</span>
<span class="pl-s1">idx</span> <span class="pl-c1">=</span> <span class="pl-s1">random</span>.<span class="pl-en">randint</span>(<span class="pl-s1">rng</span>, [<span class="pl-v">L</span><span class="pl-c1">*</span><span class="pl-v">L</span><span class="pl-c1">*</span><span class="pl-c1">8</span>], <span class="pl-s1">minval</span><span class="pl-c1">=</span><span class="pl-c1">0</span>, <span class="pl-s1">maxval</span><span class="pl-c1">=</span><span class="pl-v">L</span>)

<span class="pl-c"># Time how long it takes to compute the loss.</span>
<span class="pl-k">for</span> <span class="pl-s1">i</span> <span class="pl-c1">in</span> <span class="pl-en">range</span>(<span class="pl-c1">3</span>):
    <span class="pl-s1">start</span> <span class="pl-c1">=</span> <span class="pl-s1">time</span>.<span class="pl-en">time</span>()
    <span class="pl-s1">lv</span> <span class="pl-c1">=</span> <span class="pl-en">loss</span>(<span class="pl-s1">fvol</span>, <span class="pl-s1">idx</span>)
    <span class="pl-s1">lv</span>.<span class="pl-en">block_until_ready</span>()
    <span class="pl-s1">elapsed</span> <span class="pl-c1">=</span> <span class="pl-s1">time</span>.<span class="pl-en">time</span>() <span class="pl-c1">-</span> <span class="pl-s1">start</span>
    <span class="pl-en">print</span>(<span class="pl-s">'object data: {} - time elapsed to evaluate loss: {:.5f}'</span>.<span class="pl-en">format</span>(<span class="pl-s1">dtype</span>, <span class="pl-s1">elapsed</span>))

<span class="pl-c"># Time how long it takes to compute the gradient of the loss.</span>
<span class="pl-k">for</span> <span class="pl-s1">i</span> <span class="pl-c1">in</span> <span class="pl-en">range</span>(<span class="pl-c1">3</span>):
    <span class="pl-s1">start</span> <span class="pl-c1">=</span> <span class="pl-s1">time</span>.<span class="pl-en">time</span>()
    <span class="pl-s1">g</span> <span class="pl-c1">=</span> <span class="pl-en">loss_grad</span>(<span class="pl-s1">fvol</span>, <span class="pl-s1">idx</span>)
    <span class="pl-s1">g</span>.<span class="pl-en">block_until_ready</span>()
    <span class="pl-s1">elapsed</span> <span class="pl-c1">=</span> <span class="pl-s1">time</span>.<span class="pl-en">time</span>() <span class="pl-c1">-</span> <span class="pl-s1">start</span>
    <span class="pl-en">print</span>(<span class="pl-s">'object data: {} - time elapsed to evaluate gradient: {:.5f}'</span>.<span class="pl-en">format</span>(<span class="pl-s1">dtype</span>, <span class="pl-s1">elapsed</span>))</pre></div>
<p dir="auto">I am running this demo script on a V100 GPU with jaxlib 0.1.57 with Python 3.6.9. When I invoke the demo as <code class="notranslate">python demo.py --complex --double --index</code> (meaning that I have a complex-valued array, double precision, and a loss function that involves indexing) I get the following result:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="object data: complex128 - time elapsed to evaluate loss: 0.63968
object data: complex128 - time elapsed to evaluate loss: 0.03179
object data: complex128 - time elapsed to evaluate loss: 0.03082
object data: complex128 - time elapsed to evaluate gradient: 4.70996
object data: complex128 - time elapsed to evaluate gradient: 4.55858
object data: complex128 - time elapsed to evaluate gradient: 4.56153"><pre class="notranslate"><code class="notranslate">object data: complex128 - time elapsed to evaluate loss: 0.63968
object data: complex128 - time elapsed to evaluate loss: 0.03179
object data: complex128 - time elapsed to evaluate loss: 0.03082
object data: complex128 - time elapsed to evaluate gradient: 4.70996
object data: complex128 - time elapsed to evaluate gradient: 4.55858
object data: complex128 - time elapsed to evaluate gradient: 4.56153
</code></pre></div>
<p dir="auto">We see that it takes significantly longer to compute the gradient than the loss itself. If I invoke the script as <code class="notranslate">python demo.py --complex --double --no-index</code> (complex-valued array, double precision, and a loss that does not involve indexing) then things are better; the gradient calculation is ~3x the timing of the loss calculation.</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="object data: complex128 - time elapsed to evaluate loss: 0.18214
object data: complex128 - time elapsed to evaluate loss: 0.00098
object data: complex128 - time elapsed to evaluate loss: 0.00059
object data: complex128 - time elapsed to evaluate gradient: 0.13375
object data: complex128 - time elapsed to evaluate gradient: 0.00175
object data: complex128 - time elapsed to evaluate gradient: 0.00148"><pre class="notranslate"><code class="notranslate">object data: complex128 - time elapsed to evaluate loss: 0.18214
object data: complex128 - time elapsed to evaluate loss: 0.00098
object data: complex128 - time elapsed to evaluate loss: 0.00059
object data: complex128 - time elapsed to evaluate gradient: 0.13375
object data: complex128 - time elapsed to evaluate gradient: 0.00175
object data: complex128 - time elapsed to evaluate gradient: 0.00148
</code></pre></div>
<p dir="auto">What about in single precision? Invoking the script as <code class="notranslate">python demo.py --complex --no-double --index</code> yields</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="object data: complex64 - time elapsed to evaluate loss: 0.60020
object data: complex64 - time elapsed to evaluate loss: 0.00300
object data: complex64 - time elapsed to evaluate loss: 0.00164
object data: complex64 - time elapsed to evaluate gradient: 0.18140
object data: complex64 - time elapsed to evaluate gradient: 0.03122
object data: complex64 - time elapsed to evaluate gradient: 0.03108"><pre class="notranslate"><code class="notranslate">object data: complex64 - time elapsed to evaluate loss: 0.60020
object data: complex64 - time elapsed to evaluate loss: 0.00300
object data: complex64 - time elapsed to evaluate loss: 0.00164
object data: complex64 - time elapsed to evaluate gradient: 0.18140
object data: complex64 - time elapsed to evaluate gradient: 0.03122
object data: complex64 - time elapsed to evaluate gradient: 0.03108
</code></pre></div>
<p dir="auto">The gradient calculation is around an order of magnitude slower. What about when we have an unindexed loss function using the invocation <code class="notranslate">python demo.py --complex --no-double --no-index</code>?</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="object data: complex64 - time elapsed to evaluate loss: 0.17156
object data: complex64 - time elapsed to evaluate loss: 0.00114
object data: complex64 - time elapsed to evaluate loss: 0.00057
object data: complex64 - time elapsed to evaluate gradient: 0.11934
object data: complex64 - time elapsed to evaluate gradient: 0.00087
object data: complex64 - time elapsed to evaluate gradient: 0.00077"><pre class="notranslate"><code class="notranslate">object data: complex64 - time elapsed to evaluate loss: 0.17156
object data: complex64 - time elapsed to evaluate loss: 0.00114
object data: complex64 - time elapsed to evaluate loss: 0.00057
object data: complex64 - time elapsed to evaluate gradient: 0.11934
object data: complex64 - time elapsed to evaluate gradient: 0.00087
object data: complex64 - time elapsed to evaluate gradient: 0.00077
</code></pre></div>
<p dir="auto">Finally, what happens if I consider indexing a real-valued array instead of a complex one? Using the invocation <code class="notranslate">python demo.py --no-complex --double --index</code></p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="object data: float64 - time elapsed to evaluate loss: 0.60958
object data: float64 - time elapsed to evaluate loss: 0.00578
object data: float64 - time elapsed to evaluate loss: 0.00180
object data: float64 - time elapsed to evaluate gradient: 0.13558
object data: float64 - time elapsed to evaluate gradient: 0.00106
object data: float64 - time elapsed to evaluate gradient: 0.00082"><pre class="notranslate"><code class="notranslate">object data: float64 - time elapsed to evaluate loss: 0.60958
object data: float64 - time elapsed to evaluate loss: 0.00578
object data: float64 - time elapsed to evaluate loss: 0.00180
object data: float64 - time elapsed to evaluate gradient: 0.13558
object data: float64 - time elapsed to evaluate gradient: 0.00106
object data: float64 - time elapsed to evaluate gradient: 0.00082
</code></pre></div>
<p dir="auto">and compare this to the first output I showed where the gradient calculation was around four seconds.</p>
<p dir="auto">I note in the demo script that there is a correction to the loss that will make the code fast again when run in double precision and indexing a complex-valued array. Replacing <code class="notranslate">vv = fvol[idx].sum([-1])</code> (line 25) with</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="fvol_real = fvol.real[idx]
fvol_imag = fvol.imag[idx]
vv = jnp.sum(fvol_real + 1j*fvol_imag, [-1])"><pre class="notranslate"><span class="pl-s1">fvol_real</span> <span class="pl-c1">=</span> <span class="pl-s1">fvol</span>.<span class="pl-s1">real</span>[<span class="pl-s1">idx</span>]
<span class="pl-s1">fvol_imag</span> <span class="pl-c1">=</span> <span class="pl-s1">fvol</span>.<span class="pl-s1">imag</span>[<span class="pl-s1">idx</span>]
<span class="pl-s1">vv</span> <span class="pl-c1">=</span> <span class="pl-s1">jnp</span>.<span class="pl-en">sum</span>(<span class="pl-s1">fvol_real</span> <span class="pl-c1">+</span> <span class="pl-c1">1j</span><span class="pl-c1">*</span><span class="pl-s1">fvol_imag</span>, [<span class="pl-c1">-</span><span class="pl-c1">1</span>])</pre></div>
<p dir="auto">produces the output</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="object data: complex128 - time elapsed to evaluate loss: 0.92206
object data: complex128 - time elapsed to evaluate loss: 0.03364
object data: complex128 - time elapsed to evaluate loss: 0.03140
object data: complex128 - time elapsed to evaluate gradient: 0.20444
object data: complex128 - time elapsed to evaluate gradient: 0.06708
object data: complex128 - time elapsed to evaluate gradient: 0.06507"><pre class="notranslate"><code class="notranslate">object data: complex128 - time elapsed to evaluate loss: 0.92206
object data: complex128 - time elapsed to evaluate loss: 0.03364
object data: complex128 - time elapsed to evaluate loss: 0.03140
object data: complex128 - time elapsed to evaluate gradient: 0.20444
object data: complex128 - time elapsed to evaluate gradient: 0.06708
object data: complex128 - time elapsed to evaluate gradient: 0.06507
</code></pre></div>
<p dir="auto">which a significant improvement over the four seconds the code was taking before. This is just speculation, but in researching this issue my colleagues and I thought this bug could be at play here as well: <a class="issue-link js-issue-link" data-error-text="Failed to load title" data-id="683139267" data-permission-text="Title is private" data-url="https://github.com/google/jax/issues/4115" data-hovercard-type="issue" data-hovercard-url="/google/jax/issues/4115/hovercard" href="https://github.com/google/jax/issues/4115">#4115</a></p>
<p dir="auto">Hi,<br>
I observe a severe drop of performance of the index_update function on a GPU, when changing the data type from complex single to complex double. Is there a fundamental reason for that or is it something that could be improved?</p>
<p dir="auto">Here is a simple example to illustrate this issue:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="import jax
import jax.numpy as jnp
import numpy as np
import time

from jax.config import config
config.update(&quot;jax_enable_x64&quot;, True) # Remove this line for single precision

# Set up data
a1=jnp.zeros((1000,256),dtype=np.complex128).block_until_ready()
a2=jax.random.normal(jax.random.PRNGKey(0),(1000,256),dtype=jnp.float64) + 1.j * jax.random.normal(jax.random.PRNGKey(0),(1000,256),dtype=jnp.float64)
idx=jnp.where(jnp.abs(a2)&gt;0)
a3=a2[idx]

# Check data types used
print(a1.dtype)
print(a2.dtype)

# define jit'd index_update function
index_update_jitd=jax.jit(jax.ops.index_update)

# timing including jit
tic=time.perf_counter()
a1=index_update_jitd(a1,idx, a3).block_until_ready()
print(&quot;Execution time including jit: &quot;, time.perf_counter() -tic)

# reset data
a1=jnp.zeros((1000,256),dtype=np.complex128).block_until_ready()

# timing of jit'd function
tic=time.perf_counter()
a1=index_update_jitd(a1,idx, a3).block_until_ready()
print(&quot;Execution time of jit'd function: &quot;, time.perf_counter() -tic)"><pre class="notranslate"><code class="notranslate">import jax
import jax.numpy as jnp
import numpy as np
import time

from jax.config import config
config.update("jax_enable_x64", True) # Remove this line for single precision

# Set up data
a1=jnp.zeros((1000,256),dtype=np.complex128).block_until_ready()
a2=jax.random.normal(jax.random.PRNGKey(0),(1000,256),dtype=jnp.float64) + 1.j * jax.random.normal(jax.random.PRNGKey(0),(1000,256),dtype=jnp.float64)
idx=jnp.where(jnp.abs(a2)&gt;0)
a3=a2[idx]

# Check data types used
print(a1.dtype)
print(a2.dtype)

# define jit'd index_update function
index_update_jitd=jax.jit(jax.ops.index_update)

# timing including jit
tic=time.perf_counter()
a1=index_update_jitd(a1,idx, a3).block_until_ready()
print("Execution time including jit: ", time.perf_counter() -tic)

# reset data
a1=jnp.zeros((1000,256),dtype=np.complex128).block_until_ready()

# timing of jit'd function
tic=time.perf_counter()
a1=index_update_jitd(a1,idx, a3).block_until_ready()
print("Execution time of jit'd function: ", time.perf_counter() -tic)
</code></pre></div>
<p dir="auto">I am running this on an NVIDIA V100. Using single precision I get the output</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="complex64
complex64
Execution time including jit: 0.1222537359863054
Execution time of jit'd function: 0.0004156339855398983"><pre class="notranslate"><code class="notranslate">complex64
complex64
Execution time including jit: 0.1222537359863054
Execution time of jit'd function: 0.0004156339855398983
</code></pre></div>
<p dir="auto">With double precision it is</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="complex128
complex128
Execution time including jit: 3.707396164012607
Execution time of jit'd function: 3.5751446500071324"><pre class="notranslate"><code class="notranslate">complex128
complex128
Execution time including jit: 3.707396164012607
Execution time of jit'd function: 3.5751446500071324
</code></pre></div>
<p dir="auto">You see that for double precision the execution is almost 10000 times slower, which puzzles me.</p>
<p dir="auto">Also, the execution time seems to be roughly linear in the array size for double precision, whereas on the sizes I checked it did not really change for single precision.</p>