<p dir="auto">Please:</p>
<ul class="contains-task-list">
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox" checked=""> Check for duplicate issues. (tried my best to find a duplicate on this)</li>
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox" checked=""> Provide a complete example of how to reproduce the bug, wrapped in triple backticks like this:</li>
</ul>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="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)"><pre class="notranslate"><span class="pl-v">In</span> [<span class="pl-c1">1</span>]: <span class="pl-s1">import</span> <span class="pl-s1">jax</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-v">In</span> [<span class="pl-c1">2</span>]: <span class="pl-s1">t1</span> <span class="pl-c1">=</span> <span class="pl-s1">jnp</span>.<span class="pl-en">array</span>([[<span class="pl-c1">1.0</span>, <span class="pl-c1">2.0</span>], [<span class="pl-c1">0.0</span>, <span class="pl-c1">-</span><span class="pl-c1">1.0</span>]])
<span class="pl-v">In</span> [<span class="pl-c1">3</span>]: <span class="pl-s1">jax</span>.<span class="pl-en">grad</span>(<span class="pl-k">lambda</span> <span class="pl-s1">t</span>: <span class="pl-s1">jnp</span>.<span class="pl-en">sum</span>(<span class="pl-s1">jnp</span>.<span class="pl-en">abs</span>(<span class="pl-s1">t</span>)))(<span class="pl-s1">t1</span>)
<span class="pl-v">Out</span>[<span class="pl-c1">3</span>]: 
<span class="pl-v">DeviceArray</span>([[ <span class="pl-c1">1.</span>,  <span class="pl-c1">1.</span>],
             [ <span class="pl-c1">1.</span>, <span class="pl-c1">-</span><span class="pl-c1">1.</span>]], <span class="pl-s1">dtype</span><span class="pl-c1">=</span><span class="pl-s1">float32</span>)
<span class="pl-v">In</span> [<span class="pl-c1">7</span>]: <span class="pl-s1">t2</span> <span class="pl-c1">=</span> <span class="pl-s1">t1</span> <span class="pl-c1">*</span> <span class="pl-c1">1</span><span class="pl-c1">+</span><span class="pl-c1">0j</span> <span class="pl-c"># this is so we convert to complex64</span>
<span class="pl-v">In</span> [<span class="pl-c1">8</span>]: <span class="pl-s1">jax</span>.<span class="pl-en">grad</span>(<span class="pl-k">lambda</span> <span class="pl-s1">t</span>: <span class="pl-s1">jnp</span>.<span class="pl-en">sum</span>(<span class="pl-s1">jnp</span>.<span class="pl-en">abs</span>(<span class="pl-s1">t</span>)))(<span class="pl-s1">t2</span>)
<span class="pl-v">Out</span>[<span class="pl-c1">8</span>]: 
<span class="pl-v">DeviceArray</span>([[ <span class="pl-c1">1.</span><span class="pl-c1">+</span><span class="pl-c1">0.j</span>,  <span class="pl-c1">1.</span><span class="pl-c1">+</span><span class="pl-c1">0.j</span>],
             [ <span class="pl-c1">0.</span><span class="pl-c1">+</span><span class="pl-c1">0.j</span>, <span class="pl-c1">-</span><span class="pl-c1">1.</span><span class="pl-c1">+</span><span class="pl-c1">0.j</span>]], <span class="pl-s1">dtype</span><span class="pl-c1">=</span><span class="pl-s1">complex64</span>)</pre></div>
<p dir="auto">In the code sequence above, we can compare the grad of basically the same input as a float tensor and as a complex tensor.<br>
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 <code class="notranslate">abs</code> function applied to said entry.</p>
<p dir="auto">From looking at the source code, this seems like a matter of using <code class="notranslate">select</code> 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!</p>
<p dir="auto">Please:</p>
<ul dir="auto">
<li>[x ] Check for duplicate issues.</li>
<li>[x ] Provide a complete example of how to reproduce the bug, wrapped in triple backticks like this:</li>
</ul>
<div class="highlight highlight-source-shell notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="# clone jkonet from github
# follow jkonet setup instructions
$ python main.py --out_dir results --config_folder configs --task semicircle"><pre class="notranslate"><span class="pl-c"><span class="pl-c">#</span> clone jkonet from github</span>
<span class="pl-c"><span class="pl-c">#</span> follow jkonet setup instructions</span>
$ python main.py --out_dir results --config_folder configs --task semicircle</pre></div>
<ul dir="auto">
<li>[x ] If applicable, include full error messages/tracebacks.</li>
</ul>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="/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 &quot;/home/USER/projects/jkonet/main.py&quot;, line 292, in &lt;module&gt;
    main(args)
  File &quot;/home/USER/projects/jkonet/main.py&quot;, line 266, in main
    run_jko(config, task_dir=task_dir, logging=args.wandb)
  File &quot;/home/USER/projects/jkonet/main.py&quot;, line 32, in run_jko
    rng = jax.random.PRNGKey(int(time.time()))
  File &quot;/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/_src/random.py&quot;, line 125, in PRNGKey
    key = prng.seed_with_impl(impl, seed)
  File &quot;/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/_src/prng.py&quot;, line 233, in seed_with_impl
    return PRNGKeyArray(impl, impl.seed(seed))
  File &quot;/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/_src/prng.py&quot;, line 272, in threefry_seed
    lax.shift_right_logical(seed_arr, lax_internal._const(seed_arr, 32)))
  File &quot;/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/_src/lax/lax.py&quot;, line 487, in shift_right_logical
    return shift_right_logical_p.bind(x, y)
  File &quot;/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/core.py&quot;, line 324, in bind
    return self.bind_with_trace(find_top_trace(args), args, params)
  File &quot;/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/core.py&quot;, line 327, in bind_with_trace
    out = trace.process_primitive(self, map(trace.full_raise, args), params)
  File &quot;/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/core.py&quot;, line 684, in process_primitive
    return primitive.impl(*tracers, **params)
  File &quot;/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/_src/dispatch.py&quot;, line 99, in apply_primitive
    compiled_fun = xla_primitive_callable(prim, *unsafe_map(arg_spec, args),
  File &quot;/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/_src/util.py&quot;, line 220, in wrapper
    return cached(config._trace_context(), *args, **kwargs)
  File &quot;/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/_src/util.py&quot;, line 213, in cached
    return f(*args, **kwargs)
  File &quot;/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/_src/dispatch.py&quot;, line 164, in xla_primitive_callable
    compiled = _xla_callable_uncached(lu.wrap_init(prim_fun), device, None,
  File &quot;/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/_src/dispatch.py&quot;, line 248, in _xla_callable_uncached
    return lower_xla_callable(fun, device, backend, name, donated_invars, False,
  File &quot;/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/_src/dispatch.py&quot;, line 827, in compile
    self._executable = XlaCompiledComputation.from_xla_computation(
  File &quot;/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/_src/dispatch.py&quot;, line 934, in from_xla_computation
    compiled = compile_or_get_cached(backend, xla_computation, options,
  File &quot;/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/_src/dispatch.py&quot;, line 899, in compile_or_get_cached
    return backend_compile(backend, computation, compile_options, host_callbacks)
  File &quot;/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/_src/profiler.py&quot;, line 294, in wrapper
    return func(*args, **kwargs)
  File &quot;/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/_src/dispatch.py&quot;, 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'
"><pre class="notranslate"><code class="notranslate">/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 &lt;module&gt;
    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'

</code></pre></div>