Issue 1:
Please:

Check for duplicate issues. (tried my best to find a duplicate on this)

Provide a complete example of how to reproduce the bug, wrapped in triple backticks like this:

In [1]: import jax; import jax.numpy as jnp
In [2]: t1 = jnp.array([[1.0, 2.0], [0.0, -1.0]])
In [3]: jax.grad(lambda t: jnp.sum(jnp.abs(t)))(t1)
Out[3]:
DeviceArray([[ 1., 1.],
[ 1., -1.]], dtype=float32)
In [7]: t2 = t1 * 1+0j # this is so we convert to complex64
In [8]: jax.grad(lambda t: jnp.sum(jnp.abs(t)))(t2)
Out[8]:
DeviceArray([[ 1.+0.j, 1.+0.j],
[ 0.+0.j, -1.+0.j]], dtype=complex64)

In the code sequence above, we can compare the grad of basically the same input as a float tensor and as a complex tensor.

The problem I have (and for which I could not find any references) is that when an entry is 0, the grad changes from 1 (at float input) to 0 (at complex input), even though we're talking about the same abs function applied to said entry.

From looking at the source code, this seems like a matter of using select to replace the result at the positions which were originally zero with 1, after the calculation. I'd be happy to submit a PR if this is indeed a bug. Otherwise, perhaps a reference on why this is intended would be great!

Issue 2:
Please:

[x ] Check for duplicate issues.

[x ] Provide a complete example of how to reproduce the bug, wrapped in triple backticks like this:

# clone jkonet from github
# follow jkonet setup instructions
$ python main.py --out_dir results --config_folder configs --task semicircle

[x ] If applicable, include full error messages/tracebacks.

/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/chex/_src/pytypes.py:37: FutureWarning: jax.tree_structure is deprecated, and will be removed in a future release. Use jax.tree_util.tree_structure instead.
PyTreeDef = type(jax.tree_structure(None))
Started run.
2022-08-01 17:28:07.861579: E external/org_tensorflow/tensorflow/stream_executor/cuda/cuda_asm_compiler.cc:57] cuLinkAddData fails. This is usually caused by stale driver version.
2022-08-01 17:28:07.861611: E external/org_tensorflow/tensorflow/compiler/xla/service/gpu/gpu_compiler.cc:1325] The CUDA linking API did not work. Please use XLA_FLAGS=--xla_gpu_force_compilation_parallelism=1 to bypass it, but expect to get longer compilation time due to the lack of multi-threading.
Traceback (most recent call last):
File "/home/USER/projects/jkonet/main.py", line 292, in <module>
main(args)
File "/home/USER/projects/jkonet/main.py", line 266, in main
run_jko(config, task_dir=task_dir, logging=args.wandb)
File "/home/USER/projects/jkonet/main.py", line 32, in run_jko
rng = jax.random.PRNGKey(int(time.time()))
File "/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/_src/random.py", line 125, in PRNGKey
key = prng.seed_with_impl(impl, seed)
File "/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/_src/prng.py", line 233, in seed_with_impl
return PRNGKeyArray(impl, impl.seed(seed))
File "/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/_src/prng.py", line 272, in threefry_seed
lax.shift_right_logical(seed_arr, lax_internal._const(seed_arr, 32)))
File "/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/_src/lax/lax.py", line 487, in shift_right_logical
return shift_right_logical_p.bind(x, y)
File "/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/core.py", line 324, in bind
return self.bind_with_trace(find_top_trace(args), args, params)
File "/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/core.py", line 327, in bind_with_trace
out = trace.process_primitive(self, map(trace.full_raise, args), params)
File "/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/core.py", line 684, in process_primitive
return primitive.impl(*tracers, **params)
File "/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/_src/dispatch.py", line 99, in apply_primitive
compiled_fun = xla_primitive_callable(prim, *unsafe_map(arg_spec, args),
File "/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/_src/util.py", line 220, in wrapper
return cached(config._trace_context(), *args, **kwargs)
File "/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/_src/util.py", line 213, in cached
return f(*args, **kwargs)
File "/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/_src/dispatch.py", line 164, in xla_primitive_callable
compiled = _xla_callable_uncached(lu.wrap_init(prim_fun), device, None,
File "/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/_src/dispatch.py", line 248, in _xla_callable_uncached
return lower_xla_callable(fun, device, backend, name, donated_invars, False,
File "/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/_src/dispatch.py", line 827, in compile
self._executable = XlaCompiledComputation.from_xla_computation(
File "/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/_src/dispatch.py", line 934, in from_xla_computation
compiled = compile_or_get_cached(backend, xla_computation, options,
File "/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/_src/dispatch.py", line 899, in compile_or_get_cached
return backend_compile(backend, computation, compile_options, host_callbacks)
File "/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/_src/profiler.py", line 294, in wrapper
return func(*args, **kwargs)
File "/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/_src/dispatch.py", line 843, in backend_compile
return backend.compile(built_c, compile_options=options)
jaxlib.xla_extension.XlaRuntimeError: UNKNOWN: no kernel image is available for execution on the device
in external/org_tensorflow/tensorflow/stream_executor/cuda/cuda_asm_compiler.cc(60): 'status'