Issue 1:
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.

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("jax_enable_x64", 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))

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 python demo.py --complex --double --index (meaning that I have a complex-valued array, double precision, and a loss function that involves indexing) I get the following result:

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

We see that it takes significantly longer to compute the gradient than the loss itself. If I invoke the script as python demo.py --complex --double --no-index (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.

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

What about in single precision? Invoking the script as python demo.py --complex --no-double --index yields

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

The gradient calculation is around an order of magnitude slower. What about when we have an unindexed loss function using the invocation python demo.py --complex --no-double --no-index?

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

Finally, what happens if I consider indexing a real-valued array instead of a complex one? Using the invocation python demo.py --no-complex --double --index

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

and compare this to the first output I showed where the gradient calculation was around four seconds.

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 vv = fvol[idx].sum([-1]) (line 25) with

fvol_real = fvol.real[idx]
fvol_imag = fvol.imag[idx]
vv = jnp.sum(fvol_real + 1j*fvol_imag, [-1])

produces the output

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

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: #4115

Issue 2:
Hi,

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?

Here is a simple example to illustrate this issue:

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)>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)

I am running this on an NVIDIA V100. Using single precision I get the output

complex64
complex64
Execution time including jit: 0.1222537359863054
Execution time of jit'd function: 0.0004156339855398983

With double precision it is

complex128
complex128
Execution time including jit: 3.707396164012607
Execution time of jit'd function: 3.5751446500071324

You see that for double precision the execution is almost 10000 times slower, which puzzles me.

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.