<p dir="auto">I am getting this weird behaviour in <code class="notranslate">jnp.linalg.norm</code> where in when I set <code class="notranslate">ord="fro"</code>, I get different values when I pass input array <code class="notranslate">x</code> having shape as <code class="notranslate">[1,1]</code> instead of <code class="notranslate">[2, 1, 1]</code>. However, it is only the case for <code class="notranslate">fro</code>.<br>
For other values of <code class="notranslate">ord</code>, I get consistent values.</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.</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="import jax.numpy as jnp

x = jnp.zeros([2, 1, 1])
x = x.at[0, 0, 0].set(3.4346476e-38)
x
# DeviceArray([[[3.4346476e-38]],
#
#             [[0.0000000e+00]]], dtype=float32)

y = jnp.zeros([1, 1])
y = y.at[0, 0].set(3.4346476e-38)
y
# DeviceArray([[3.4346476e-38]], dtype=float32)

print(jnp.linalg.norm(x, &quot;fro&quot;, (-2, -1)))
print(jnp.linalg.norm(y, &quot;fro&quot;, (-2, -1)))
# [3.4346476e-38 0.0000000e+00]
# 0.0 -------&gt; why is this zero ? and not 3.4346476e-38 ?

print(jnp.linalg.norm(x, 1, (-2, -1)))
print(jnp.linalg.norm(y, 1, (-2, -1)))
# [3.4346476e-38 0.0000000e+00]
# 3.4346476e-38

print(jnp.linalg.norm(x, 2, (-2, -1)))
print(jnp.linalg.norm(y, 2, (-2, -1)))
# [3.4346476e-38 0.0000000e+00]
# 3.4346476e-38"><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">x</span> <span class="pl-c1">=</span> <span class="pl-s1">jnp</span>.<span class="pl-en">zeros</span>([<span class="pl-c1">2</span>, <span class="pl-c1">1</span>, <span class="pl-c1">1</span>])
<span class="pl-s1">x</span> <span class="pl-c1">=</span> <span class="pl-s1">x</span>.<span class="pl-s1">at</span>[<span class="pl-c1">0</span>, <span class="pl-c1">0</span>, <span class="pl-c1">0</span>].<span class="pl-en">set</span>(<span class="pl-c1">3.4346476e-38</span>)
<span class="pl-s1">x</span>
<span class="pl-c"># DeviceArray([[[3.4346476e-38]],</span>
<span class="pl-c">#</span>
<span class="pl-c">#             [[0.0000000e+00]]], dtype=float32)</span>

<span class="pl-s1">y</span> <span class="pl-c1">=</span> <span class="pl-s1">jnp</span>.<span class="pl-en">zeros</span>([<span class="pl-c1">1</span>, <span class="pl-c1">1</span>])
<span class="pl-s1">y</span> <span class="pl-c1">=</span> <span class="pl-s1">y</span>.<span class="pl-s1">at</span>[<span class="pl-c1">0</span>, <span class="pl-c1">0</span>].<span class="pl-en">set</span>(<span class="pl-c1">3.4346476e-38</span>)
<span class="pl-s1">y</span>
<span class="pl-c"># DeviceArray([[3.4346476e-38]], dtype=float32)</span>

<span class="pl-en">print</span>(<span class="pl-s1">jnp</span>.<span class="pl-s1">linalg</span>.<span class="pl-en">norm</span>(<span class="pl-s1">x</span>, <span class="pl-s">"fro"</span>, (<span class="pl-c1">-</span><span class="pl-c1">2</span>, <span class="pl-c1">-</span><span class="pl-c1">1</span>)))
<span class="pl-en">print</span>(<span class="pl-s1">jnp</span>.<span class="pl-s1">linalg</span>.<span class="pl-en">norm</span>(<span class="pl-s1">y</span>, <span class="pl-s">"fro"</span>, (<span class="pl-c1">-</span><span class="pl-c1">2</span>, <span class="pl-c1">-</span><span class="pl-c1">1</span>)))
<span class="pl-c"># [3.4346476e-38 0.0000000e+00]</span>
<span class="pl-c"># 0.0 -------&gt; why is this zero ? and not 3.4346476e-38 ?</span>

<span class="pl-en">print</span>(<span class="pl-s1">jnp</span>.<span class="pl-s1">linalg</span>.<span class="pl-en">norm</span>(<span class="pl-s1">x</span>, <span class="pl-c1">1</span>, (<span class="pl-c1">-</span><span class="pl-c1">2</span>, <span class="pl-c1">-</span><span class="pl-c1">1</span>)))
<span class="pl-en">print</span>(<span class="pl-s1">jnp</span>.<span class="pl-s1">linalg</span>.<span class="pl-en">norm</span>(<span class="pl-s1">y</span>, <span class="pl-c1">1</span>, (<span class="pl-c1">-</span><span class="pl-c1">2</span>, <span class="pl-c1">-</span><span class="pl-c1">1</span>)))
<span class="pl-c"># [3.4346476e-38 0.0000000e+00]</span>
<span class="pl-c"># 3.4346476e-38</span>

<span class="pl-en">print</span>(<span class="pl-s1">jnp</span>.<span class="pl-s1">linalg</span>.<span class="pl-en">norm</span>(<span class="pl-s1">x</span>, <span class="pl-c1">2</span>, (<span class="pl-c1">-</span><span class="pl-c1">2</span>, <span class="pl-c1">-</span><span class="pl-c1">1</span>)))
<span class="pl-en">print</span>(<span class="pl-s1">jnp</span>.<span class="pl-s1">linalg</span>.<span class="pl-en">norm</span>(<span class="pl-s1">y</span>, <span class="pl-c1">2</span>, (<span class="pl-c1">-</span><span class="pl-c1">2</span>, <span class="pl-c1">-</span><span class="pl-c1">1</span>)))
<span class="pl-c"># [3.4346476e-38 0.0000000e+00]</span>
<span class="pl-c"># 3.4346476e-38</span></pre></div>
<ul class="contains-task-list">
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox" checked=""> If applicable, include full error messages/tracebacks.</li>
</ul>
<p dir="auto">Simple example:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="jnp.linalg.norm(jnp.array([1e-20], jnp.float32))"><pre class="notranslate"><code class="notranslate">jnp.linalg.norm(jnp.array([1e-20], jnp.float32))
</code></pre></div>
<p dir="auto">produces 0.</p>
<p dir="auto">Denormals are flushed to zero, and 1e-20**2 is smaller than the smallest normal float32.</p>
<p dir="auto">It's possible to define a non-underflowing Euclidean norm using variadic reductions:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="def norm2(xs):
  def reducer(xs, ys):
    (ssq_x, scale_x) = xs
    (ssq_y, scale_y) = ys
    scale = lax.max(scale_x, scale_y)
    scale_is_zero = lax.eq(scale, jnp.float32(0))
    ssq = ssq_x * (scale_x/scale)** 2 + ssq_y * (scale_y / scale)**2
    return (lax.select(scale_is_zero, jnp.ones_like(ssq), ssq), scale)


  ssq_y, scale_y = lax.reduce(
    (np.ones_like(x), jnp.abs(x)),
    (jnp.float32(1), jnp.float32(0)),
    reducer, dimensions=(0,))
  return scale_y * jnp.sqrt(ssq_y)"><pre class="notranslate"><code class="notranslate">def norm2(xs):
  def reducer(xs, ys):
    (ssq_x, scale_x) = xs
    (ssq_y, scale_y) = ys
    scale = lax.max(scale_x, scale_y)
    scale_is_zero = lax.eq(scale, jnp.float32(0))
    ssq = ssq_x * (scale_x/scale)** 2 + ssq_y * (scale_y / scale)**2
    return (lax.select(scale_is_zero, jnp.ones_like(ssq), ssq), scale)


  ssq_y, scale_y = lax.reduce(
    (np.ones_like(x), jnp.abs(x)),
    (jnp.float32(1), jnp.float32(0)),
    reducer, dimensions=(0,))
  return scale_y * jnp.sqrt(ssq_y)
</code></pre></div>
<p dir="auto">We should determine whether we want to do this; if so it probably needs to be implemented as a primitive.</p>
<p dir="auto">Interestingly <code class="notranslate">numpy.linalg.norm</code> also produces 0 if the CPU <code class="notranslate">ftz</code> flag is set (e.g., via the <code class="notranslate">daz</code> package), but <code class="notranslate">scipy.linalg.norm</code> does not.</p>