<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 requests.</li>
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox" checked=""> Describe your goal, and if possible provide a code snippet with a motivating example.</li>
</ul>
<p dir="auto">Using 16-bit (half-precision) floating point numbers reduces memory usage and can make the program run faster, compared to f.ex. 32-bit. My program is working with complex numbers and I'd like to take advantage of the 16-bit speed-up, but JAX only support <code class="notranslate">complex64</code>.</p>
<p dir="auto">Can we add support for half-precision complex numbers (<code class="notranslate">complex32</code>)?</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="import jax.numpy as jnp

f32 = jnp.array([1, 2, 3], dtype=&quot;float32&quot;)
print(f32.size * f32.itemsize)  # 12 (bytes)

f16 = jnp.array([1, 2, 3], dtype=&quot;bfloat16&quot;)
print(f16.size * f16.itemsize)  # 6 (bytes)

c64 = jnp.array([1, 2, 3], dtype=&quot;complex64&quot;)
print(c64.size * c64.itemsize)  # 24 (bytes)

c32 = jnp.array([1, 2, 3], dtype=&quot;complex32&quot;)
print(c32.size * c32.itemsize)  # 12 (bytes)  &lt;- I want this, but complex32 is not supported"><pre class="notranslate"><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-s1">f32</span> <span class="pl-c1">=</span> <span class="pl-s1">jnp</span>.<span class="pl-en">array</span>([<span class="pl-c1">1</span>, <span class="pl-c1">2</span>, <span class="pl-c1">3</span>], <span class="pl-s1">dtype</span><span class="pl-c1">=</span><span class="pl-s">"float32"</span>)
<span class="pl-en">print</span>(<span class="pl-s1">f32</span>.<span class="pl-s1">size</span> <span class="pl-c1">*</span> <span class="pl-s1">f32</span>.<span class="pl-s1">itemsize</span>)  <span class="pl-c"># 12 (bytes)</span>

<span class="pl-s1">f16</span> <span class="pl-c1">=</span> <span class="pl-s1">jnp</span>.<span class="pl-en">array</span>([<span class="pl-c1">1</span>, <span class="pl-c1">2</span>, <span class="pl-c1">3</span>], <span class="pl-s1">dtype</span><span class="pl-c1">=</span><span class="pl-s">"bfloat16"</span>)
<span class="pl-en">print</span>(<span class="pl-s1">f16</span>.<span class="pl-s1">size</span> <span class="pl-c1">*</span> <span class="pl-s1">f16</span>.<span class="pl-s1">itemsize</span>)  <span class="pl-c"># 6 (bytes)</span>

<span class="pl-s1">c64</span> <span class="pl-c1">=</span> <span class="pl-s1">jnp</span>.<span class="pl-en">array</span>([<span class="pl-c1">1</span>, <span class="pl-c1">2</span>, <span class="pl-c1">3</span>], <span class="pl-s1">dtype</span><span class="pl-c1">=</span><span class="pl-s">"complex64"</span>)
<span class="pl-en">print</span>(<span class="pl-s1">c64</span>.<span class="pl-s1">size</span> <span class="pl-c1">*</span> <span class="pl-s1">c64</span>.<span class="pl-s1">itemsize</span>)  <span class="pl-c"># 24 (bytes)</span>

<span class="pl-s1">c32</span> <span class="pl-c1">=</span> <span class="pl-s1">jnp</span>.<span class="pl-en">array</span>([<span class="pl-c1">1</span>, <span class="pl-c1">2</span>, <span class="pl-c1">3</span>], <span class="pl-s1">dtype</span><span class="pl-c1">=</span><span class="pl-s">"complex32"</span>)
<span class="pl-en">print</span>(<span class="pl-s1">c32</span>.<span class="pl-s1">size</span> <span class="pl-c1">*</span> <span class="pl-s1">c32</span>.<span class="pl-s1">itemsize</span>)  <span class="pl-c"># 12 (bytes)  &lt;- I want this, but complex32 is not supported</span></pre></div>
<p dir="auto">If not, is there anything that stops me from making JAX work with half-precision complex numbers today (f.ex. by creating a class that has two fields: <code class="notranslate">real &lt;bfloat16&gt;</code> and <code class="notranslate">imag &lt;bfloat16&gt;</code>)?</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="from typing import NamedTuple
import jax.numpy as jnp

class Complex32(NamedTuple):
    real: jnp.ndarray
    imag: jnp.ndarray

c32 = Complex32(
    jnp.array([1, 2, 3], dtype=&quot;bfloat16&quot;),
    jnp.array([4, 5, 6], dtype=&quot;bfloat16&quot;),
)
jnp.abs(c32)  # TypeError: absolute requires ndarray or scalar arguments, got &lt;class '__main__.Complex32'&gt; at position 0."><pre class="notranslate"><span class="pl-k">from</span> <span class="pl-s1">typing</span> <span class="pl-k">import</span> <span class="pl-v">NamedTuple</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">class</span> <span class="pl-v">Complex32</span>(<span class="pl-v">NamedTuple</span>):
    <span class="pl-s1">real</span>: <span class="pl-s1">jnp</span>.<span class="pl-s1">ndarray</span>
    <span class="pl-s1">imag</span>: <span class="pl-s1">jnp</span>.<span class="pl-s1">ndarray</span>

<span class="pl-s1">c32</span> <span class="pl-c1">=</span> <span class="pl-v">Complex32</span>(
    <span class="pl-s1">jnp</span>.<span class="pl-en">array</span>([<span class="pl-c1">1</span>, <span class="pl-c1">2</span>, <span class="pl-c1">3</span>], <span class="pl-s1">dtype</span><span class="pl-c1">=</span><span class="pl-s">"bfloat16"</span>),
    <span class="pl-s1">jnp</span>.<span class="pl-en">array</span>([<span class="pl-c1">4</span>, <span class="pl-c1">5</span>, <span class="pl-c1">6</span>], <span class="pl-s1">dtype</span><span class="pl-c1">=</span><span class="pl-s">"bfloat16"</span>),
)
<span class="pl-s1">jnp</span>.<span class="pl-en">abs</span>(<span class="pl-s1">c32</span>)  <span class="pl-c"># TypeError: absolute requires ndarray or scalar arguments, got &lt;class '__main__.Complex32'&gt; at position 0.</span></pre></div>
<p dir="auto">Related issue: <a class="issue-link js-issue-link" data-error-text="Failed to load title" data-id="397819991" data-permission-text="Title is private" data-url="https://github.com/google/jax/issues/216" data-hovercard-type="issue" data-hovercard-url="/google/jax/issues/216/hovercard" href="https://github.com/google/jax/issues/216">#216</a></p>
<p dir="auto">Thanks!</p>
<p dir="auto"><a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/hawkinsp/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/hawkinsp">@hawkinsp</a> when generating HLO with many reductions, JAX duplicates reduction computations. Despite that not affecting codegen much, the resulting HLO dumps get overly huge and hard to work with due to such computations. Similarly to tf2xla it would be very difficult if JAX could deduplicate such computations.<br>
<a target="_blank" rel="noopener noreferrer nofollow" href="https://user-images.githubusercontent.com/348959/129822274-a74cb97d-474f-41d4-98ee-c2829c7c3d87.png"><img src="https://user-images.githubusercontent.com/348959/129822274-a74cb97d-474f-41d4-98ee-c2829c7c3d87.png" alt="unnamed" style="max-width: 100%;"></a></p>