Unable to export quantized models to stablehlo

Hello, I am currently trying out torchax, specifically the torchax.export.exported_program_to_stablehlo. While trying out its capabilities, I ran into a problem when trying to export quantized models via the function. As far as I can see it seems that this is simply not supported.

So I wanted to ask if this is not supported yet or a bug. In case of it being a bug, I provide the following code to reproduce my errors:

```
import torch
import torchax.export
import torch.nn as nn
from torchao.quantization import quantize_
from torchao.quantization.quant_api import Int4WeightOnlyConfig

def llama_mlp_example():
class LlamaMLP(nn.Module):
def __init__(self):
super().__init__()
self.hidden_size = 16
self.intermediate_size = 64
self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
self.act_fn = nn.SiLU()
def forward(self, x):
down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
return down_proj
m = LlamaMLP().eval()
example_inputs = (torch.empty(1, 16, dtype=torch.float32),)
return m, example_inputs

def simple_pt2e_quant_example(m, example_inputs):
m = torch.export.export(m, example_inputs).module()

from torchao.quantization.pt2e.quantize_pt2e import (
prepare_pt2e,
convert_pt2e,
)

from executorch.backends.xnnpack.quantizer.xnnpack_quantizer import (
get_symmetric_quantization_config,
XNNPACKQuantizer,
)

quantizer = XNNPACKQuantizer().set_global(get_symmetric_quantization_config())
m = prepare_pt2e(m, quantizer)

# calibration omitted

m = convert_pt2e(m)

exported = torch.export.export(m, example_inputs)

weights, stablehlo = torchax.export.exported_program_to_stablehlo(exported)

def int4_quant_example(m, example_inputs):
quant_config = Int4WeightOnlyConfig(
group_size=32,
)

quantize_(m, quant_config)

exported = torch.export.export(
m,
example_inputs,
)

weights, stablehlo = torchax.export.exported_program_to_stablehlo(exported)

if __name__ == "__main__":
llama_mlp_m, llama_mlp_inputs = llama_mlp_example()

simple_pt2e_quant_example(llama_mlp_m, llama_mlp_inputs)
# int4_quant_example(llama_mlp_m, llama_mlp_inputs)

```
The errors that occur are as follows:

```
AssertionError: quantized_decomposed.dequantize_per_tensor.default

While executing %dequantize_per_tensor : [num_users=1] = call_function[target=torch.ops.quantized_decomposed.dequantize_per_tensor.default](args = (%b__frozen_param0, 1.0, 0, -127, 127, torch.int8), kwargs = {})
Original traceback:
File "/home/Ha-Eyz/testing/issueExample.py", line 19, in forward
down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
File "/home/Ha-Eyz/roof-mlir/.venv/lib/python3.11/site-packages/torch/nn/modules/linear.py", line 134, in forward
return F.linear(input, self.weight, self.bias)
Use tlparse to see full graph. (https://github.com/pytorch/tlparse?tab=readme-ov-file#tlparse-parse-structured-pt2-logs)
--------------------
For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.
```

If you run the simple_pt2e_quant_example(...) function

```
AssertionError: fbgemm.bf16i4bf16_rowwise.default

While executing %bf16i4bf16_rowwise : [num_users=1] = call_function[target=torch.ops.fbgemm.bf16i4bf16_rowwise.default](args = (%view, %p_down_proj_parametrizations_weight_original0, %p_down_proj_parametrizations_weight_original1, %p_down_proj_parametrizations_weight_original2), kwargs = {})
Original traceback:
File "/home/Ha-Eyz/testing/issueExample.py", line 19, in forward
down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
File "/home/Ha-Eyz/roof-mlir/.venv/lib/python3.11/site-packages/torch/nn/modules/linear.py", line 134, in forward
return F.linear(input, self.weight, self.bias)
Use tlparse to see full graph. (https://github.com/pytorch/tlparse?tab=readme-ov-file#tlparse-parse-structured-pt2-logs)
--------------------
For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.
```

If you run the int4_quant_example(...) function

For the full error messages you can run the code I provided above.