<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">I think h100 support fp8. And it is actually inside cuda12.1. Is there ant plan to support this datatype?</p>
<p dir="auto">So <code class="notranslate">vmap</code> looks awesome and I'd like to use it everywhere I possibly can. However it looks like it's not really compatible with stax, and in particular convolutions.</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="from functools import partial

import jax.numpy as np
from jax import random, vmap
from jax.experimental.stax import Conv, GeneralConv

net_init, net_apply = Conv(32, (3, 3), padding='SAME')

in_shape = (28, 28, 1)
out_shape, net_params = net_init(random.PRNGKey(0), in_shape)

# Apply network to dummy inputs
inputs = np.zeros((128, 28, 28, 1))
predictions = vmap(partial(net_apply, net_params))(inputs)
print(predictions.shape)"><pre class="notranslate"><span class="pl-k">from</span> <span class="pl-s1">functools</span> <span class="pl-k">import</span> <span class="pl-s1">partial</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">np</span>
<span class="pl-k">from</span> <span class="pl-s1">jax</span> <span class="pl-k">import</span> <span class="pl-s1">random</span>, <span class="pl-s1">vmap</span>
<span class="pl-k">from</span> <span class="pl-s1">jax</span>.<span class="pl-s1">experimental</span>.<span class="pl-s1">stax</span> <span class="pl-k">import</span> <span class="pl-v">Conv</span>, <span class="pl-v">GeneralConv</span>

<span class="pl-s1">net_init</span>, <span class="pl-s1">net_apply</span> <span class="pl-c1">=</span> <span class="pl-v">Conv</span>(<span class="pl-c1">32</span>, (<span class="pl-c1">3</span>, <span class="pl-c1">3</span>), <span class="pl-s1">padding</span><span class="pl-c1">=</span><span class="pl-s">'SAME'</span>)

<span class="pl-s1">in_shape</span> <span class="pl-c1">=</span> (<span class="pl-c1">28</span>, <span class="pl-c1">28</span>, <span class="pl-c1">1</span>)
<span class="pl-s1">out_shape</span>, <span class="pl-s1">net_params</span> <span class="pl-c1">=</span> <span class="pl-en">net_init</span>(<span class="pl-s1">random</span>.<span class="pl-v">PRNGKey</span>(<span class="pl-c1">0</span>), <span class="pl-s1">in_shape</span>)

<span class="pl-c"># Apply network to dummy inputs</span>
<span class="pl-s1">inputs</span> <span class="pl-c1">=</span> <span class="pl-s1">np</span>.<span class="pl-en">zeros</span>((<span class="pl-c1">128</span>, <span class="pl-c1">28</span>, <span class="pl-c1">28</span>, <span class="pl-c1">1</span>))
<span class="pl-s1">predictions</span> <span class="pl-c1">=</span> <span class="pl-en">vmap</span>(<span class="pl-en">partial</span>(<span class="pl-s1">net_apply</span>, <span class="pl-s1">net_params</span>))(<span class="pl-s1">inputs</span>)
<span class="pl-en">print</span>(<span class="pl-s1">predictions</span>.<span class="pl-s1">shape</span>)</pre></div>
<p dir="auto">gives an error</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="---------------------------------------------------------------------------
IndexError                                Traceback (most recent call last)
~/nu/skainswo/research/gan_with_the_wind/vae.py in &lt;module&gt;
     13 
     14 in_shape = (28, 28, 1)
---&gt; 15 out_shape, net_params = net_init(random.PRNGKey(0), in_shape)
     16 
     17 # Apply network to dummy inputs

~/.local/share/virtualenvs/research-OGGq2tNy/lib/python3.7/site-packages/jax/experimental/stax.py in init_fun(rng, input_shape)
    112     kernel_shape = [out_chan if c == 'O' else
    113                     input_shape[lhs_spec.index('C')] if c == 'I' else
--&gt; 114                     next(filter_shape_iter) for c in rhs_spec]
    115     output_shape = lax.conv_general_shape_tuple(
    116         input_shape, kernel_shape, strides, padding, dimension_numbers)

~/.local/share/virtualenvs/research-OGGq2tNy/lib/python3.7/site-packages/jax/experimental/stax.py in &lt;listcomp&gt;(.0)
    112     kernel_shape = [out_chan if c == 'O' else
    113                     input_shape[lhs_spec.index('C')] if c == 'I' else
--&gt; 114                     next(filter_shape_iter) for c in rhs_spec]
    115     output_shape = lax.conv_general_shape_tuple(
    116         input_shape, kernel_shape, strides, padding, dimension_numbers)

IndexError: tuple index out of range"><pre class="notranslate"><code class="notranslate">---------------------------------------------------------------------------
IndexError                                Traceback (most recent call last)
~/nu/skainswo/research/gan_with_the_wind/vae.py in &lt;module&gt;
     13 
     14 in_shape = (28, 28, 1)
---&gt; 15 out_shape, net_params = net_init(random.PRNGKey(0), in_shape)
     16 
     17 # Apply network to dummy inputs

~/.local/share/virtualenvs/research-OGGq2tNy/lib/python3.7/site-packages/jax/experimental/stax.py in init_fun(rng, input_shape)
    112     kernel_shape = [out_chan if c == 'O' else
    113                     input_shape[lhs_spec.index('C')] if c == 'I' else
--&gt; 114                     next(filter_shape_iter) for c in rhs_spec]
    115     output_shape = lax.conv_general_shape_tuple(
    116         input_shape, kernel_shape, strides, padding, dimension_numbers)

~/.local/share/virtualenvs/research-OGGq2tNy/lib/python3.7/site-packages/jax/experimental/stax.py in &lt;listcomp&gt;(.0)
    112     kernel_shape = [out_chan if c == 'O' else
    113                     input_shape[lhs_spec.index('C')] if c == 'I' else
--&gt; 114                     next(filter_shape_iter) for c in rhs_spec]
    115     output_shape = lax.conv_general_shape_tuple(
    116         input_shape, kernel_shape, strides, padding, dimension_numbers)

IndexError: tuple index out of range
</code></pre></div>
<p dir="auto">which makes some sense based on the <code class="notranslate">dimension_numbers</code> in the definition here: <a href="https://github.com/google/jax/blob/master/jax/experimental/stax.py#L127">https://github.com/google/jax/blob/master/jax/experimental/stax.py#L127</a>. But using</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="net_init, net_apply = GeneralConv((&quot;HWC&quot;, &quot;HWIO&quot;, &quot;HWC&quot;), 32, (3, 3))"><pre class="notranslate"><span class="pl-s1">net_init</span>, <span class="pl-s1">net_apply</span> <span class="pl-c1">=</span> <span class="pl-v">GeneralConv</span>((<span class="pl-s">"HWC"</span>, <span class="pl-s">"HWIO"</span>, <span class="pl-s">"HWC"</span>), <span class="pl-c1">32</span>, (<span class="pl-c1">3</span>, <span class="pl-c1">3</span>))</pre></div>
<p dir="auto">doesn't work either:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
~/nu/skainswo/research/gan_with_the_wind/vae.py in &lt;module&gt;
     13 
     14 in_shape = (28, 28, 1)
---&gt; 15 out_shape, net_params = net_init(random.PRNGKey(0), in_shape)
     16 
     17 # Apply network to dummy inputs

~/.local/share/virtualenvs/research-OGGq2tNy/lib/python3.7/site-packages/jax/experimental/stax.py in init_fun(rng, input_shape)
    114                     next(filter_shape_iter) for c in rhs_spec]
    115     output_shape = lax.conv_general_shape_tuple(
--&gt; 116         input_shape, kernel_shape, strides, padding, dimension_numbers)
    117     bias_shape = [out_chan if c == 'C' else 1 for c in out_spec]
    118     bias_shape = tuple(itertools.dropwhile(lambda x: x == 1, bias_shape))

~/.local/share/virtualenvs/research-OGGq2tNy/lib/python3.7/site-packages/jax/lax/lax.py in conv_general_shape_tuple(lhs_shape, rhs_shape, window_strides, padding, dimension_numbers)
   3858 def conv_general_shape_tuple(lhs_shape, rhs_shape, window_strides, padding,
   3859                              dimension_numbers):
-&gt; 3860   lhs_perm, rhs_perm, out_perm = conv_general_permutations(dimension_numbers)
   3861   lhs_trans = onp.take(lhs_shape, lhs_perm)
   3862   rhs_trans = onp.take(rhs_shape, rhs_perm)

~/.local/share/virtualenvs/research-OGGq2tNy/lib/python3.7/site-packages/jax/lax/lax.py in conv_general_permutations(dimension_numbers)
   4002       msg = (&quot;convolution dimension_numbers[{}] must contain the characters &quot;
   4003              &quot;'{}' and '{}' exatly once, got {}.&quot;)
-&gt; 4004       raise TypeError(msg.format(i, a, b, dimension_numbers[i]))
   4005     if len(dimension_numbers[i]) != len(set(dimension_numbers[i])):
   4006       msg = (&quot;convolution dimension_numbers[{}] cannot have duplicate &quot;

TypeError: convolution dimension_numbers[0] must contain the characters 'N' and 'C' exatly once, got HWC."><pre class="notranslate"><code class="notranslate">---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
~/nu/skainswo/research/gan_with_the_wind/vae.py in &lt;module&gt;
     13 
     14 in_shape = (28, 28, 1)
---&gt; 15 out_shape, net_params = net_init(random.PRNGKey(0), in_shape)
     16 
     17 # Apply network to dummy inputs

~/.local/share/virtualenvs/research-OGGq2tNy/lib/python3.7/site-packages/jax/experimental/stax.py in init_fun(rng, input_shape)
    114                     next(filter_shape_iter) for c in rhs_spec]
    115     output_shape = lax.conv_general_shape_tuple(
--&gt; 116         input_shape, kernel_shape, strides, padding, dimension_numbers)
    117     bias_shape = [out_chan if c == 'C' else 1 for c in out_spec]
    118     bias_shape = tuple(itertools.dropwhile(lambda x: x == 1, bias_shape))

~/.local/share/virtualenvs/research-OGGq2tNy/lib/python3.7/site-packages/jax/lax/lax.py in conv_general_shape_tuple(lhs_shape, rhs_shape, window_strides, padding, dimension_numbers)
   3858 def conv_general_shape_tuple(lhs_shape, rhs_shape, window_strides, padding,
   3859                              dimension_numbers):
-&gt; 3860   lhs_perm, rhs_perm, out_perm = conv_general_permutations(dimension_numbers)
   3861   lhs_trans = onp.take(lhs_shape, lhs_perm)
   3862   rhs_trans = onp.take(rhs_shape, rhs_perm)

~/.local/share/virtualenvs/research-OGGq2tNy/lib/python3.7/site-packages/jax/lax/lax.py in conv_general_permutations(dimension_numbers)
   4002       msg = ("convolution dimension_numbers[{}] must contain the characters "
   4003              "'{}' and '{}' exatly once, got {}.")
-&gt; 4004       raise TypeError(msg.format(i, a, b, dimension_numbers[i]))
   4005     if len(dimension_numbers[i]) != len(set(dimension_numbers[i])):
   4006       msg = ("convolution dimension_numbers[{}] cannot have duplicate "

TypeError: convolution dimension_numbers[0] must contain the characters 'N' and 'C' exatly once, got HWC.
</code></pre></div>