Integer convolution with int8 inputs and int32 output

[vidulae]

I'm trying to do a convolution with an int8 input image and kernel tensors, then get back an int32 output tensor, but I don't see any way to do that:

kernel = jax.numpy.arange(3*3*3*64).reshape(3,3,3,64).astype(jax.numpy.int8)
img = jax.numpy.arange(32*1024*1024*3).reshape(32,1024,1024,3).astype(jax.numpy.int8)
out = jax.lax.conv(lhs=jax.numpy.transpose(img,[0,3,1,2]), rhs=jax.numpy.transpose(kernel,[3,2,0,1]), window_strides=(1, 1), padding='SAME')

gives back an int8 result. But there is an option to give preferred_element_type to get the output from a preferred type. But it doesn't support int32 or int16 types. It only supports float32 and int8 types.

out = jax.lax.conv(lhs=jax.numpy.transpose(img,[0,3,1,2]), rhs=jax.numpy.transpose(kernel,[3,2,0,1]), window_strides=(1, 1), padding='SAME', preferred_element_type ='jax.numpy.float32')

I could cast the output tensor to jax.numpy.int32 after the convolution, but I want to use a hardware-accelerated int8 convolution, if available. Are there any available solutions in order to do that?

Thanks!