Issue 1:
Please:

Check for duplicate requests.

Describe your goal, and if possible provide a code snippet with a motivating example.

I think h100 support fp8. And it is actually inside cuda12.1. Is there ant plan to support this datatype?

Issue 2:
So vmap 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.

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)

gives an error

---------------------------------------------------------------------------
IndexError Traceback (most recent call last)
~/nu/skainswo/research/gan_with_the_wind/vae.py in <module>
13
14 in_shape = (28, 28, 1)
---> 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
--> 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 <listcomp>(.0)
112 kernel_shape = [out_chan if c == 'O' else
113 input_shape[lhs_spec.index('C')] if c == 'I' else
--> 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

which makes some sense based on the dimension_numbers in the definition here: https://github.com/google/jax/blob/master/jax/experimental/stax.py#L127. But using

net_init, net_apply = GeneralConv(("HWC", "HWIO", "HWC"), 32, (3, 3))

doesn't work either:

---------------------------------------------------------------------------
TypeError Traceback (most recent call last)
~/nu/skainswo/research/gan_with_the_wind/vae.py in <module>
13
14 in_shape = (28, 28, 1)
---> 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(
--> 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):
-> 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 {}.")
-> 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.