int4 type support for matmul

[rybakov]

Hello,

I am working on hardware-accelerated matmul with int4 input tensors and get back an int16/int32 tensor.
I would like to be able to run matmul and emit MHLO with example below (or let me know if there are better options):

x=jnp.array(100).reshape((1,1)).astype(jnp.int4)
jnp.matmul(x,x)

Thanks!

[zhangqiaorjc]

cc @hawkinsp

[hawkinsp]

This is fixed these days!

In [5]: import jax, jax.numpy as jnp, numpy as np
In [8]: x=np.array(100).reshape((1,1)).astype(jnp.int4)
   ...: print(jax.jit(jnp.matmul).lower(x,x).as_text())
module @jit_matmul attributes {mhlo.num_partitions = 1 : i32, mhlo.num_replicas = 1 : i32} {
  func.func public @main(%arg0: tensor<1x1xi4> {mhlo.sharding = "{replicated}"}, %arg1: tensor<1x1xi4> {mhlo.sharding = "{replicated}"}) -> (tensor<1x1xi4> {jax.result_info = ""}) {
    %0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0], precision = [DEFAULT, DEFAULT] : (tensor<1x1xi4>, tensor<1x1xi4>) -> tensor<1x1xi4>
    return %0 : tensor<1x1xi4>
  }
}

(There's no guarantee any given JAX backend knows how to compile that, but producing the stablehlo is no problem.)