diff --git a/qwix/contrib/hijax/hiqarray.py b/qwix/contrib/hijax/hiqarray.py
new file mode 100644
index 0000000..766eb7b
--- /dev/null
+++ b/qwix/contrib/hijax/hiqarray.py
@@ -0,0 +1,523 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Assumes that group size = -1 means full tensor quantization along that axis.
+
+We assume that the cotangent type of a QArray is an Array.
+We could always add a function that tells us how to do the backward stuff.
+For example, if we want to quantize the backwards on a matmul we could do that.
+"""
+
+import dataclasses
+
+import jax
+import jax.experimental.hijax as hjx
+import jax.experimental.pallas as pl
+import jax.numpy as jnp
+import numpy as np
+import qwix.contrib.hijax.hiqarray_common as hqc
+
+TransformedRef = jax._src.state.types.TransformedRef  # pylint: disable=protected-access
+NDIndexer = jax._src.state.indexing.NDIndexer  # pylint: disable=protected-access
+
+
+@dataclasses.dataclass(frozen=True)
+class HiQArray:
+  """HiQArray class."""
+
+  # Arrays
+  qvalue: jax.Array
+  scale: jax.Array
+  zero_point: jax.Array | None
+
+  # Quantization metadata
+  metadata: hqc.QuantizationMetadata
+
+  @property
+  def shape(self):
+    return self.qvalue.shape
+
+  @property
+  def dtype(self):
+    return self.scale.dtype
+
+  @property
+  def qtype(self):
+    return self.qvalue.dtype
+
+
+# Generic functions for QArrays
+# These are utility functions and should not be called outside of this file
+def _qarray_lo_ty(
+    qmd: hqc.QuantizationMetadata, use_zero_point: bool
+) -> list[hjx.ShapedArray]:
+  """Returns the lo_ty of a QArray."""
+  out = [
+      hjx.ShapedArray(qmd.data_shape, qmd.qtype),
+      hjx.ShapedArray(qmd.quant_shape, qmd.dtype),
+  ]
+  if use_zero_point:
+    out.append(hjx.ShapedArray(qmd.quant_shape, qmd.qtype))
+  return out
+
+
+def _qarray_lower_val(hi_val: HiQArray) -> list[jax.Array]:
+  """Returns the lo_val of a QArray."""
+  if hi_val.zero_point is None:
+    return [hi_val.qvalue, hi_val.scale]
+  else:
+    return [hi_val.qvalue, hi_val.scale, hi_val.zero_point]
+
+
+def _qarray_lower_block_spec(
+    metadata: hqc.QuantizationMetadata,
+    use_zero_point: bool,
+    block_spec: pl.BlockSpec,
+) -> (
+    tuple[pl.BlockSpec, pl.BlockSpec]
+    | tuple[pl.BlockSpec, pl.BlockSpec, pl.BlockSpec]
+):
+  """Returns the block spec of a QArray."""
+  data_block_spec = block_spec
+  block_spec_size = tuple(
+      map(lambda b: getattr(b, "block_size", b), block_spec.block_shape)  # pyrefly: ignore
+  )  # pyrefly: ignore
+  scale_block_shapes = hqc.map_ints_over_shapes(
+      block_spec_size, metadata.data_shape, metadata.quant_shape  # pyrefly: ignore
+  )  # pyrefly: ignore
+
+  scale_block_spec = pl.BlockSpec(scale_block_shapes, block_spec.index_map)
+
+  if use_zero_point:
+    zero_point_block_spec = scale_block_spec
+    return data_block_spec, scale_block_spec, zero_point_block_spec
+  else:
+    return data_block_spec, scale_block_spec
+
+
+# QArray type
+@dataclasses.dataclass(frozen=True)
+class HiQArrayTy(hjx.HiType):
+  """HiQArray type."""
+
+  metadata: hqc.QuantizationMetadata
+  use_zero_point: bool
+  shape = property(lambda self: self.metadata.data_shape)
+  dtype = property(lambda self: self.metadata.dtype)
+
+  def update(self, *, shape: tuple[int, ...] | None = None, **kwargs):
+    # Used for refs
+    if len(kwargs) > 0:
+      assert False, f"Unsupported kwargs: {kwargs=}"
+
+    if shape is None:
+      return self
+    if len(shape) == len(self.metadata.data_shape) and all(
+        s == q for s, q in zip(shape, self.metadata.data_shape)
+    ):
+      return self
+
+    # Compute new metadata
+    scale_shape = hqc.map_ints_over_shapes(
+        shape, self.shape, self.metadata.quant_shape
+    )
+    data_shape_dtype_struct = jax.ShapeDtypeStruct(shape, self.metadata.qtype)
+    scale_shape_dtype_struct = jax.ShapeDtypeStruct(
+        scale_shape, self.metadata.dtype
+    )
+    new_qmd = hqc.QuantizationMetadata.init_from_qvalue_and_scales(
+        data_shape_dtype_struct, scale_shape_dtype_struct  # pyrefly: ignore
+    )  # pyrefly: ignore
+
+    return HiQArrayTy(new_qmd, self.use_zero_point)
+
+  def lo_ty(self) -> list[hjx.ShapedArray]:  # pyrefly: ignore
+    return _qarray_lo_ty(self.metadata, self.use_zero_point)  # pyrefly: ignore
+
+  # Functions for raising to hijax and lowering to lojax
+  def lower_val(self, hi_val: HiQArray) -> list[jax.Array]:
+    return _qarray_lower_val(hi_val)
+
+  def raise_val(  # pyrefly: ignore
+      self, data: jax.Array, scale: jax.Array, zero_point: jax.Array | None
+  ) -> HiQArray:
+    return HiQArray(data, scale, zero_point, self.metadata)  # pyrefly: ignore
+
+  # Autodiff functions
+  def to_tangent_aval(self):
+    return hjx.ShapedArray(self.metadata.data_shape, self.metadata.dtype)
+
+  def to_ct_aval(self):
+    return hjx.ShapedArray(self.metadata.data_shape, self.metadata.dtype)
+
+  def vspace_zero(self):
+    return jnp.zeros(self.metadata.data_shape, self.metadata.dtype)
+
+  # Lowering for block specs
+  def lower_block_spec(self, block_spec: pl.BlockSpec):
+    return _qarray_lower_block_spec(
+        self.metadata, self.use_zero_point, block_spec
+    )
+
+  # Ref handling
+  def ref_get_abstract_eval(self, ref_aval, *args, tree):
+    arr_aval = jax.core.ShapedArray(self.shape, self.metadata.dtype)
+    updated_ref = ref_aval.update(inner_aval=arr_aval)
+    out, effects = jax._src.state.primitives.get_p.abstract_eval(  # pylint: disable=protected-access
+        updated_ref, *args, tree=tree
+    )
+    assert isinstance(out, jax.core.ShapedArray)
+
+    # Do some math to figure out the scale shapes
+    scale_shape = hqc.map_ints_over_shapes(
+        out.shape, self.shape, self.metadata.quant_shape
+    )
+    data_shape_dtype_struct = jax.ShapeDtypeStruct(
+        out.shape, self.metadata.qtype
+    )
+    scale_shape_dtype_struct = jax.ShapeDtypeStruct(
+        scale_shape, self.metadata.dtype
+    )
+    new_qmd = hqc.QuantizationMetadata.init_from_qvalue_and_scales(
+        data_shape_dtype_struct, scale_shape_dtype_struct  # pyrefly: ignore
+    )  # pyrefly: ignore
+    return HiQArrayTy(new_qmd, self.use_zero_point), effects
+
+  def ref_swap_abstract_eval(self, ref_aval, val_aval, *args, tree):
+    arr_aval = jax.core.ShapedArray(self.shape, self.metadata.dtype)
+    val_arr_aval = jax.core.ShapedArray(val_aval.shape, self.metadata.dtype)
+    updated_ref = ref_aval.update(inner_aval=arr_aval)
+    out_aval, effects = jax._src.state.primitives.swap_p.abstract_eval(  # pylint: disable=protected-access
+        updated_ref, val_arr_aval, *args, tree=tree
+    )
+    assert isinstance(out_aval, jax.core.ShapedArray)
+
+    # Do some math to figure out the scale shapes
+    scale_shape = hqc.map_ints_over_shapes(
+        out_aval.shape, self.shape, self.metadata.quant_shape
+    )
+    data_shape_dtype_struct = jax.ShapeDtypeStruct(
+        out_aval.shape, self.metadata.qtype
+    )
+    scale_shape_dtype_struct = jax.ShapeDtypeStruct(
+        scale_shape, self.metadata.dtype
+    )
+    new_qmd = hqc.QuantizationMetadata.init_from_qvalue_and_scales(
+        data_shape_dtype_struct, scale_shape_dtype_struct  # pyrefly: ignore
+    )  # pyrefly: ignore
+    return HiQArrayTy(new_qmd, self.use_zero_point), effects
+
+  def ref_get_to_lojax(self, ref: TransformedRef | jax.Ref, idx: NDIndexer):
+
+    if isinstance(ref, TransformedRef):
+      if ref.transforms:
+        raise NotImplementedError(ref)
+      ref = ref.ref
+    # Unpack Ref type
+    unpacked_ref = ref._refs  # pylint: disable=protected-access
+
+    data_slices = idx.indices
+    scale_slices = hqc.map_slices_over_shapes(
+        idx.indices, self.shape, self.metadata.quant_shape  # pyrefly: ignore
+    )  # pyrefly: ignore
+    lovals = self.lower_val(unpacked_ref)
+    slices = [data_slices, scale_slices, scale_slices]
+
+    # If zero_point isn't used, then the zip won't hit it
+    outs = [out.get()[*slcs] for out, slcs in zip(lovals, slices)]  # pyrefly: ignore
+    return self.raise_val(*outs)
+
+  def ref_swap_to_lojax(
+      self, ref: TransformedRef | jax.Ref, val: jax.Array, idx: NDIndexer
+  ):
+
+    if isinstance(ref, TransformedRef):
+      if ref.transforms:
+        raise NotImplementedError(ref)
+      ref = ref.ref
+    # Unpack Ref type
+    unpacked_ref = ref._refs  # pylint: disable=protected-access
+
+    data_slices = idx.indices
+    scale_slices = hqc.map_slices_over_shapes(
+        idx.indices, self.shape, self.metadata.quant_shape  # pyrefly: ignore
+    )  # pyrefly: ignore
+    slices = [data_slices, scale_slices, scale_slices]
+
+    outs = [
+        out.swap(val)[*slcs]  # pytype: disable=attribute-error
+        for out, val, slcs in zip(
+            self.lower_val(unpacked_ref), self.lower_val(val), slices  # pytype: disable=wrong-arg-types
+        )
+    ]
+    return self.raise_val(*outs)
+
+  def __repr__(self):
+    return f"HiQArrayTy({self.metadata.jaxpr_repr(self.use_zero_point)})"
+
+
+hjx.register_hitype(
+    HiQArray, lambda q: HiQArrayTy(q.metadata, q.zero_point is not None)
+)
+
+
+class ToHiQArray(hjx.VJPHiPrimitive):
+  """Hijax primitive to convert data, scale, and zero_point to a HiQArray."""
+
+  def __init__(
+      self,
+      metadata: hqc.QuantizationMetadata,
+      key: jax.Array | None,
+      use_zero_point: bool,
+      use_lower: bool,
+      use_upper: bool,
+  ):
+    data_scale_avals = (
+        hjx.ShapedArray(metadata.data_shape, metadata.dtype),
+        hjx.ShapedArray(metadata.quant_shape, metadata.dtype),
+    )
+    zp_aval = (
+        hjx.ShapedArray(metadata.quant_shape, metadata.qtype)
+        if use_zero_point
+        else None
+    )
+    lower_aval = (
+        hjx.ShapedArray(metadata.quant_shape, metadata.dtype)
+        if use_lower
+        else None
+    )
+    upper_aval = (
+        hjx.ShapedArray(metadata.quant_shape, metadata.dtype)
+        if use_upper
+        else None
+    )
+    self.in_avals = (
+        data_scale_avals + (zp_aval,) + (lower_aval,) + (upper_aval,)
+    )
+    self.out_aval = HiQArrayTy(metadata, use_zero_point)
+    self.params = dict(
+        metadata=metadata,
+        key=key,
+        use_zero_point=use_zero_point,
+        use_lower=use_lower,
+        use_upper=use_upper,
+    )
+    # For pytype warnings - redundant since super() will initialize these.
+    self.metadata = metadata
+    self.key = key
+    self.use_zero_point = use_zero_point
+    self.use_lower = use_lower
+    self.use_upper = use_upper
+
+    super().__init__()
+
+  def expand(self, data, scale, zero_point, lower, upper):  # pyrefly: ignore
+    quantized_data = hqc.scale_and_round(
+        data,
+        scale,
+        zero_point,
+        self.metadata,
+        key=self.key,
+        lower=lower,
+        upper=upper,
+    )
+    return HiQArray(quantized_data, scale, zero_point, self.metadata)
+
+  # Reverse mode ad
+  def vjp_fwd(self, nzs_in, data, scale, zero_point, lower, upper):  # pyrefly: ignore
+    return to_hiqarray(
+        data,
+        scale,
+        zero_point,
+        self.metadata,
+        key=self.key,
+        lower=lower,
+        upper=upper,
+    ), (data, scale, zero_point, lower, upper)
+
+  def vjp_bwd_retval(self, res, outgrad, /):
+    # Classic API: returns values instead of using accumulators
+    fn = lambda *vals: hqc.scale_and_round(
+        vals[0],
+        vals[1],
+        vals[2],
+        self.metadata,
+        lower=vals[3],
+        upper=vals[4],
+        key=self.key,
+        differentiable=True,
+    )
+    _, vjp_fn = jax.vjp(fn, *res)
+    vjp = vjp_fn(outgrad)
+    return vjp
+
+  def transpose(self, ct, *maybe_accums):
+    pass
+
+  def batch_dim_rule(self, axis_data, dims, /):
+    raise NotImplementedError(
+        f"for vmap support, subclass {type(self)} must "
+        "implement `batch` or `batch_dim_rule`"
+    )
+
+
+def to_hiqarray(
+    data: jax.Array,
+    scale: jax.Array,
+    zero_point: jax.Array | None,
+    metadata: hqc.QuantizationMetadata,
+    key: jax.Array | None = None,
+    lower=None,
+    upper=None,
+) -> HiQArray:
+  """Converts data, scale, and zero_point to a HiQArray."""
+  to_qarray_instance = ToHiQArray(
+      metadata,
+      key,
+      zero_point is not None,
+      lower is not None,
+      upper is not None,
+  )
+  return to_qarray_instance(data, scale, zero_point, lower, upper)
+
+
+class FromHiQArray(hjx.VJPHiPrimitive):
+  """Hijax primitive to convert a HiQArray to a dequantized array."""
+
+  def __init__(self, metadata: hqc.QuantizationMetadata, use_zero_point: bool):
+    self.in_avals = (HiQArrayTy(metadata, use_zero_point),)
+    self.out_aval = hjx.ShapedArray(metadata.data_shape, metadata.dtype)
+    self.params = dict(metadata=metadata, use_zero_point=use_zero_point)
+    # For pytype warnings - redundant since super() will initialize these.
+    self.metadata = metadata
+    self.use_zero_point = use_zero_point
+    super().__init__()
+
+  def expand(self, qarray: HiQArray):  # pyrefly: ignore
+    dequantized_data = hqc.scale_and_round_inverse(  # pyrefly: ignore
+        qarray.qvalue, qarray.scale, qarray.zero_point, self.metadata
+    )
+    return dequantized_data
+
+  # Reverse mode ad
+  def vjp_fwd(self, nzs_in, qarray: HiQArray):  # pyrefly: ignore
+    return from_hiqarray(qarray), (
+        qarray.qvalue,
+        qarray.scale,
+        qarray.zero_point,
+    )  # pyrefly: ignore
+
+  def vjp_bwd_retval(self, res, outgrad, /):
+    # Classic API: returns values instead of using accumulators
+    fn = lambda x: hqc.scale_and_round_inverse(x, res[1], res[2], self.metadata)
+    _, vjp_fn = jax.vjp(fn, res[0].astype(self.metadata.dtype))
+    vjp = vjp_fn(outgrad)
+    return vjp
+
+  def transpose(self, ct, *maybe_accums):
+    pass
+
+  def batch_dim_rule(self, axis_data, dims, /):
+    raise NotImplementedError(
+        f"for vmap support, subclass {type(self)} must "
+        "implement `batch` or `batch_dim_rule`"
+    )
+
+
+def from_hiqarray(qarray: HiQArray) -> jax.Array:
+  """Converts a HiQArray to a dequantized array."""
+  ty = jax.typeof(qarray)
+  from_qarray_instance = FromHiQArray(qarray.metadata, ty.use_zero_point)
+  return from_qarray_instance(qarray)
+
+
+class PermuteAxes(hjx.VJPHiPrimitive):
+  """Hijax primitive to permute the axes of a HiQArray."""
+
+  def __init__(
+      self,
+      metadata: hqc.QuantizationMetadata,
+      use_zero_point: bool,
+      axes: tuple[int, ...],
+  ):
+    self.in_avals = (HiQArrayTy(metadata, use_zero_point),)
+    self.out_aval = HiQArrayTy(
+        PermuteAxes._permute_axes_metadata(metadata, axes), use_zero_point
+    )
+    self.params = dict(
+        metadata=metadata,
+        use_zero_point=use_zero_point,
+        axes=axes,
+    )
+    # For pytype warnings - redundant since super() will initialize these.
+    self.metadata = metadata
+    self.use_zero_point = use_zero_point
+    self.axes = axes
+    super().__init__()
+
+  # Private functions
+  @staticmethod
+  def _permute_axes_metadata(
+      metadata: hqc.QuantizationMetadata, axes: tuple[int, ...]
+  ) -> hqc.QuantizationMetadata:
+    in_data_shape_dtype_struct = jax.ShapeDtypeStruct(
+        metadata.data_shape, metadata.qtype
+    )
+    in_scale_shape_dtype_struct = jax.ShapeDtypeStruct(
+        metadata.quant_shape, metadata.dtype
+    )
+    new_data_shape_dtype_struct = jax.eval_shape(
+        lambda x: jnp.permute_dims(x, axes), in_data_shape_dtype_struct
+    )
+    new_scale_shape_dtype_struct = jax.eval_shape(
+        lambda x: jnp.permute_dims(x, axes), in_scale_shape_dtype_struct
+    )
+    new_qmd = hqc.QuantizationMetadata.init_from_qvalue_and_scales(
+        new_data_shape_dtype_struct, new_scale_shape_dtype_struct
+    )
+    return new_qmd
+
+  def expand(self, qarray: HiQArray):  # pyrefly: ignore
+    return HiQArray(
+        jnp.permute_dims(qarray.qvalue, self.axes),
+        jnp.permute_dims(qarray.scale, self.axes),
+        jnp.permute_dims(qarray.zero_point, self.axes)
+        if qarray.zero_point is not None
+        else None,
+        self.metadata,
+    )  # pyrefly: ignore
+
+  # Reverse mode ad
+  def vjp_fwd(self, nzs_in, qarray: HiQArray):  # pyrefly: ignore
+    return permute_axes(qarray, self.axes), tuple()  # pyrefly: ignore
+
+  def vjp_bwd_retval(self, res, outgrad, /):
+    # Classic API: returns values instead of using accumulators
+    inv_perm = tuple(np.argsort(self.axes))
+    return (jnp.permute_dims(outgrad, inv_perm),)
+
+  def transpose(self, ct, *maybe_accums):
+    pass
+
+  def batch_dim_rule(self, axis_data, dims, /):
+    raise NotImplementedError(
+        f"for vmap support, subclass {type(self)} must "
+        "implement `batch` or `batch_dim_rule`"
+    )
+
+
+def permute_axes(qarray: HiQArray, axes: tuple[int, ...]) -> HiQArray:
+  """Permutes the axes of a HiQArray."""
+  ty = jax.typeof(qarray)
+  permute_axes_instance = PermuteAxes(qarray.metadata, ty.use_zero_point, axes)
+  return permute_axes_instance(qarray)
diff --git a/qwix/contrib/hijax/hiqarray_common.py b/qwix/contrib/hijax/hiqarray_common.py
index 759de79..7fa57e8 100644
--- a/qwix/contrib/hijax/hiqarray_common.py
+++ b/qwix/contrib/hijax/hiqarray_common.py
@@ -17,8 +17,13 @@
 """
 
 import dataclasses
+from typing import overload
 
+import jax
 import jax.numpy as jnp
+import qwix.contrib.hijax.hiquant_utils as hq_utils
+
+JSlice = jax._src.indexing.Slice  # pylint: disable=protected-access
 
 
 @dataclasses.dataclass(frozen=True)
@@ -54,16 +59,16 @@ def init(
       cls,
       data_shape: tuple[int, ...],
       quant_info: dict[int, int],
-      original_dtype: jnp.dtype,
-      quantized_dtype: jnp.dtype,
+      dtype: jnp.dtype,
+      qtype: jnp.dtype,
   ):
     """Initializes the quantization metadata for an array.
 
     Args:
       data_shape: The shape of the original array.
       quant_info: A dictionary of quantization axes to group sizes.
-      original_dtype: The dtype of the original array.
-      quantized_dtype: The dtype of the quantized array.
+      dtype: The dtype of the original array.
+      qtype: The dtype of the quantized array.
 
     Returns:
       The quantization metadata for the array.
@@ -100,8 +105,30 @@ def init(
         quant_compatible_shape,
         tiled_reduction_axes,
         full_reduction_axes,
-        original_dtype,
-        quantized_dtype,
+        dtype,
+        qtype,
+    )
+
+  @classmethod
+  def init_from_qvalue_and_scales(cls, qvalue: jax.Array, scales: jax.Array):
+    """Initializes the quantization metadata for an array from the quantized value and scales."""
+    dtype = scales.dtype
+    qtype = qvalue.dtype
+    data_shape = qvalue.shape
+    quant_axes = tuple([
+        i for i, (x, y) in enumerate(zip(qvalue.shape, scales.shape)) if x != y
+    ])
+    group_sizes = tuple([
+        x // y
+        for i, (x, y) in enumerate(zip(qvalue.shape, scales.shape))
+        if x != y
+    ])
+    quant_info = {k: v for k, v in zip(quant_axes, group_sizes)}
+    return cls.init(
+        data_shape,
+        quant_info,
+        dtype,
+        qtype,
     )
 
   @staticmethod
@@ -212,7 +239,7 @@ def _get_quant_shape(
         out.append(xi)
     return tuple(out)
 
-  def __repr__(self):
+  def detailed_repr(self):
     quant_axes = self.quant_axes
     group_sizes = self.group_sizes
     orig_dtype = self.dtype
@@ -222,3 +249,156 @@ def __repr__(self):
         f" {quant_dtype=})"
     )
     return out
+
+  def __repr__(self):
+    quant_type_str = str(jax.core.ShapedArray(self.quant_shape, self.qtype))
+
+    logical_type = jax.core.ShapedArray(self.data_shape, self.dtype)
+
+    # find the type information
+    i = quant_type_str.find("[")
+    qtype_str = quant_type_str[:i]
+    quant_shape_str = quant_type_str[i:]
+
+    out = (
+        f"QMD(logical_type={logical_type}, qtype={qtype_str},"
+        f" qshape={quant_shape_str})"
+    )
+    return out
+
+  def jaxpr_repr(self, use_zero_point: bool):
+    data_type = jax.core.ShapedArray(self.data_shape, self.qtype)
+    scale_type = jax.core.ShapedArray(self.quant_shape, self.dtype)
+    zero_point_type = jax.core.ShapedArray(self.quant_shape, self.qtype)
+    if use_zero_point:
+      out = f"{data_type}, {scale_type}, {zero_point_type}"
+    else:
+      out = f"{data_type}, {scale_type}, None"
+    return out
+
+
+# Functions that operate on Arrays
+
+
+def scale_and_round(
+    data: jax.Array,
+    scale: jax.Array,
+    zero_point: jax.Array | None,
+    metadata: QuantizationMetadata,
+    *,
+    lower: float | None = None,
+    upper: float | None = None,
+    key: jax.Array | None = None,
+    differentiable: bool = False,
+) -> jax.Array:
+  """Quantization operation."""
+  # e.g. q = clip(round(x/scale + zero_point), lower, upper)
+  start_shape = data.shape
+  uqd = data.reshape(metadata.data_compatible_shape)
+  sc = scale.reshape(metadata.quant_compatible_shape)
+  if zero_point is None:
+    out = uqd / sc
+  else:
+    zp = zero_point.reshape(metadata.quant_compatible_shape)
+    out = uqd / sc + zp
+
+  # Stochastic rounding
+  if key is not None:
+    out += jax.random.uniform(key, shape=out.shape, minval=-0.5, maxval=0.5)
+
+  out = jnp.clip(out, lower, upper)
+
+  if not differentiable:
+    if hq_utils.is_integer_dtype(metadata.qtype):
+      out = jnp.round(out)
+    out = out.astype(metadata.qtype)
+  out = out.reshape(start_shape)
+
+  return out
+
+
+def scale_and_round_inverse(
+    data: jax.Array,
+    scale: jax.Array,
+    zero_point: jax.Array | None,
+    metadata: QuantizationMetadata,
+) -> jax.Array:
+  """Basic dequantization method."""
+  # e.g. x = scale * (q - zero_point)
+  qd = data.reshape(metadata.data_compatible_shape)
+  sc = scale.reshape(metadata.quant_compatible_shape)
+  if zero_point is None:
+    out = sc * qd.astype(metadata.dtype)
+  else:
+    zp = zero_point.reshape(metadata.quant_compatible_shape)
+    out = sc * (qd.astype(metadata.dtype) - zp)
+  return out.reshape(metadata.data_shape)
+
+
+@overload
+def map_slice(slc: slice, big_shape: int, small_shape: int) -> slice:
+  ...
+
+
+@overload
+def map_slice(slc: JSlice, big_shape: int, small_shape: int) -> JSlice:
+  ...
+
+
+def map_slice(
+    slc: slice | JSlice, big_shape: int, small_shape: int
+) -> slice | JSlice:
+  """Maps a slice from a big shape to a small shape."""
+
+  use_jax_slice = isinstance(slc, JSlice)
+
+  if use_jax_slice:
+    slc = slice(slc.start, slc.start + slc.size, slc.stride)
+  assert slc.step is None or slc.step == 1, "Only unit strides supported"
+
+  # We assume that big_shape is a multiple of small_shape
+  ratio = big_shape // small_shape
+
+  num_big_elements = slc.stop - slc.start
+  if num_big_elements % ratio != 0:
+    raise ValueError("Slice size must be a multiple of ratio")
+  num_small_elements = num_big_elements // ratio
+
+  new_start = slc.start // ratio
+  new_stop = new_start + num_small_elements
+
+  if use_jax_slice:
+    return JSlice(new_start, new_stop, 1)
+  return slice(new_start, new_stop, 1)
+
+
+def map_int(i: int, big_shape: int, small_shape: int) -> int:
+  """Maps an int from a big shape to a small shape."""
+  ratio = big_shape // small_shape
+  if i % ratio != 0:
+    raise ValueError("Index must be a multiple of ratio")
+  return i // ratio
+
+
+def map_slices_over_shapes(
+    slcs: tuple[slice | JSlice, ...],
+    big_shapes: tuple[int, ...],
+    small_shapes: tuple[int, ...],
+) -> tuple[slice | JSlice, ...]:
+  """Maps a slice from a big shape to a small shape."""
+  assert len(big_shapes) == len(
+      small_shapes
+  ), "Big and small shapes must have same length"
+  return tuple(map(map_slice, slcs, big_shapes, small_shapes))  # pyrefly: ignore
+
+
+def map_ints_over_shapes(
+    ints: tuple[int, ...],
+    big_shapes: tuple[int, ...],
+    small_shapes: tuple[int, ...],
+) -> tuple[int, ...]:
+  """Maps an int from a big shape to a small shape."""
+  assert len(big_shapes) == len(
+      small_shapes
+  ), "Big and small shapes must have same length"
+  return tuple(map(map_int, ints, big_shapes, small_shapes))  # pyrefly: ignore
diff --git a/qwix/contrib/hijax/test_utils.py b/qwix/contrib/hijax/test_utils.py
new file mode 100644
index 0000000..cf02b7e
--- /dev/null
+++ b/qwix/contrib/hijax/test_utils.py
@@ -0,0 +1,54 @@
+"""Utility functions for testing hiqarrays."""
+
+import jax
+import jax.numpy as jnp
+from qwix.contrib.hijax import hiqarray as hq
+from qwix.contrib.hijax import hiqarray_common as hqc
+
+
+def create_random_qarray(
+    orig_shape: tuple[int, ...],
+    quant_axes: tuple[int, ...],
+    group_sizes: tuple[int, ...],
+    key: jax.Array | None = None,
+) -> hq.HiQArray:
+  """Create a HiQArray with random values."""
+  if key is None:
+    key = jax.random.key(0)
+  k1, k2, k3 = jax.random.split(key, 3)
+  data = jax.random.randint(
+      k1, orig_shape, minval=-128, maxval=128, dtype=jnp.int8
+  )
+  quant_info = {k: v for k, v in zip(quant_axes, group_sizes)}
+  metadata = hqc.QuantizationMetadata.init(
+      orig_shape, quant_info, jnp.float32, jnp.int8
+  )
+  scale = jax.random.normal(k2, metadata.quant_shape, dtype=jnp.float32)
+  zero_point = jax.random.randint(
+      k3, metadata.quant_shape, minval=-128, maxval=128, dtype=jnp.int8
+  )
+  return hq.HiQArray(data, scale, zero_point, metadata)
+
+
+def create_invertible_qarray(
+    orig_shape: tuple[int, ...],
+    quant_axes: tuple[int, ...],
+    group_sizes: tuple[int, ...],
+    use_zero_point: bool,
+    key: jax.Array | None = None,
+) -> hq.HiQArray:
+  """Create a HiQArray that can be inverted exactly with to_qarray and from_qarray."""
+  if key is None:
+    key = jax.random.key(0)
+  data = jax.random.normal(key, orig_shape, dtype=jnp.float32)
+  quant_info = {k: v for k, v in zip(quant_axes, group_sizes)}
+  metadata = hqc.QuantizationMetadata.init(
+      orig_shape, quant_info, jnp.float32, jnp.float32
+  )
+  scale = jnp.ones(metadata.quant_shape, dtype=jnp.float32)
+  zero_point = (
+      jnp.zeros(metadata.quant_shape, dtype=jnp.float32)
+      if use_zero_point
+      else None
+  )
+  return hq.HiQArray(data, scale, zero_point, metadata)
diff --git a/tests/contrib/hijax/hiqarray_test.py b/tests/contrib/hijax/hiqarray_test.py
new file mode 100644
index 0000000..0f323ed
--- /dev/null
+++ b/tests/contrib/hijax/hiqarray_test.py
@@ -0,0 +1,345 @@
+from absl.testing import absltest
+import jax
+import jax.experimental.pallas as pl
+import jax.experimental.pallas.tpu as pltpu
+import jax.numpy as jnp
+import numpy as np
+from qwix.contrib.hijax import test_utils
+import qwix.contrib.hijax.hiqarray as hq
+import qwix.contrib.hijax.hiqarray_common as hqc
+
+
+class HiQArrayTest(absltest.TestCase):
+
+  def test_qarray_exists(self):
+    qarray = test_utils.create_random_qarray((16, 16), (0, 1), (16, 16))
+    self.assertIsInstance(qarray, hq.HiQArray)
+
+  # Test other transforms and basic functionality here
+  def test_ref(self):
+    qarray = test_utils.create_random_qarray((16, 16), (0, 1), (16, 16))
+    _ = jax.new_ref(qarray)
+
+  def test_ref_get(self):
+    qarray = test_utils.create_random_qarray((16, 16), (0, 1), (1, 1))
+
+    @jax.jit
+    def f(q):
+      ref = jax.new_ref(q)
+      return ref[:, 0:2]
+
+    o = f(qarray)
+    assert jnp.all(o.qvalue == qarray.qvalue[:, 0:2])
+    assert jnp.all(o.scale == qarray.scale[:, 0:2])
+
+    assert o.shape == (16, 2)
+    assert o.metadata.data_shape == (16, 2)
+    assert o.metadata.quant_shape == (16, 2)
+
+  def test_ref_swap(self):
+    qarray1 = test_utils.create_random_qarray((16, 16), (0, 1), (16, 16))
+    qarray2 = test_utils.create_random_qarray((16, 16), (0, 1), (16, 16))
+
+    @jax.jit
+    def f(q1, q2):
+      ref = jax.new_ref(q1)
+      ref[:, :] = q2
+      return ref.get()
+
+    o = f(qarray1, qarray2)
+    assert jnp.all(o.qvalue == qarray2.qvalue)
+    assert jnp.all(o.scale == qarray2.scale)
+
+  def test_ref_swap_2(self):
+    qarray1 = test_utils.create_random_qarray((8, 8), (0, 1), (1, 1))
+    qarray2 = test_utils.create_random_qarray((16, 16), (0, 1), (1, 1))
+
+    @jax.jit
+    def f(q1, q2):
+      ref = jax.new_ref(q1)
+      ref2 = jax.new_ref(q2)
+      ref[:, :] = ref2[:8, :8]
+      return ref.get()
+
+    o = f(qarray1, qarray2)
+    assert jnp.all(o.qvalue == qarray2.qvalue[:8, :8])
+    assert jnp.all(o.scale == qarray2.scale[:8, :8])
+
+    assert o.shape == (8, 8)
+    assert o.metadata.data_shape == (8, 8)
+    assert o.metadata.quant_shape == (8, 8)
+
+  def test_simple_kernel(self):
+    qarray = test_utils.create_random_qarray((16, 16), (0, 1), (16, 16))
+
+    def fn(q):
+      def q_kernel(q_ref, o_ref):
+        o_ref[...] = hq.from_hiqarray(q_ref[...])
+
+      return pl.pallas_call(
+          q_kernel,
+          out_shape=jax.ShapeDtypeStruct(q.shape, q.dtype),
+          interpret=pltpu.InterpretParams(),
+      )(q)
+
+    fn(qarray)
+
+  def test_bad_block_spec(self):
+    qarray = test_utils.create_random_qarray((16, 16), (0, 1), (16, 16))
+
+    block_spec = pl.BlockSpec((2, 2), lambda i, j: (i, j))
+    ty = jax.typeof(qarray)
+    with self.assertRaises(ValueError):
+      ty.lower_block_spec(block_spec)
+
+  def test_pallas_call(self):
+    qarray = test_utils.create_random_qarray((16, 16), (0, 1), (1, 1))
+
+    def fn(q, *, bm=8, bn=8):
+      m, n = q.shape
+
+      def q_kernel(q_ref, o_ref):
+        assert q_ref.shape == (bm, bn)
+        out_arr = hq.from_hiqarray(q_ref[...])
+        o_ref[...] = out_arr
+
+      out_shape = jax.ShapeDtypeStruct(q.shape, q.dtype)
+      out_spec = pl.BlockSpec((bm, bn), lambda i, j: (i, j))
+      grid = (m // bm, n // bn)
+      return pl.pallas_call(
+          q_kernel,
+          out_shape=out_shape,
+          out_specs=out_spec,
+          grid=grid,
+          in_specs=[pl.BlockSpec((bm, bn), lambda i, j: (i, j))],
+          interpret=pltpu.InterpretParams(),
+      )(q)
+
+    fn(qarray)
+
+
+class ToHiQArrayTest(absltest.TestCase):
+
+  def setUp(self):
+    super().setUp()
+    k1, k2, k3 = jax.random.split(jax.random.key(0), 3)
+    orig_shape = (16, 16)
+    quant_info = {0: 2, 1: 2}
+    self.qvalue = jax.random.normal(k1, orig_shape)
+    self.metadata = hqc.QuantizationMetadata.init(
+        orig_shape, quant_info, jnp.float32, jnp.int8
+    )
+    self.scale = jax.random.normal(
+        k2, self.metadata.quant_shape, dtype=jnp.float32
+    )
+    self.zero_point = jax.random.normal(
+        k3, self.metadata.quant_shape, dtype=jnp.float32
+    )
+
+  def test_to_qarray_with_zero_point(self):
+    # Checks that to_qarray works
+    qarray = hq.to_hiqarray(
+        self.qvalue,
+        self.scale,
+        self.zero_point,
+        self.metadata,
+        key=jax.random.key(0),
+    )
+    self.assertIsInstance(qarray, hq.HiQArray)
+
+  def test_to_qarray_without_zero_point(self):
+    qarray = hq.to_hiqarray(
+        self.qvalue, self.scale, None, self.metadata, key=jax.random.key(0)
+    )
+    self.assertIsInstance(qarray, hq.HiQArray)
+
+  def test_to_qarray_backward_with_zero_point(self):
+    # Set up simple examples for autograd testing
+    scale = jnp.ones_like(self.scale)
+    zp = jnp.zeros_like(self.zero_point)
+
+    fn = lambda data, scale, zero_point: hq.to_hiqarray(
+        data, scale, zero_point, self.metadata, key=jax.random.key(0)
+    )
+    qarray = hq.to_hiqarray(
+        self.qvalue, scale, zp, self.metadata, key=jax.random.key(0)
+    )
+
+    primals_out, bwd_fn = jax.vjp(fn, self.qvalue, scale, zp)
+
+    # Test that primals are correct
+    np.testing.assert_allclose(primals_out.qvalue, qarray.qvalue)
+
+    # Compute cotangents
+    cotangent = jax.random.normal(
+        jax.random.key(0), self.qvalue.shape, dtype=self.qvalue.dtype
+    )
+    next_cotangent = bwd_fn(cotangent)
+
+    # Check that cotangents are correct
+    np.testing.assert_allclose(next_cotangent[0], cotangent)
+
+  def test_to_qarray_backward_without_zero_point(self):
+    # Set up simple examples for autograd testing
+    scale = jnp.ones_like(self.scale)
+
+    fn = lambda data, scale: hq.to_hiqarray(
+        data, scale, None, self.metadata, key=jax.random.key(0)
+    )
+    qarray = hq.to_hiqarray(
+        self.qvalue, scale, None, self.metadata, key=jax.random.key(0)
+    )
+
+    primals_out, bwd_fn = jax.vjp(fn, self.qvalue, scale)
+
+    # Test that primals are correct
+    np.testing.assert_allclose(primals_out.qvalue, qarray.qvalue)
+
+    # Compute cotangents
+    cotangent = jax.random.normal(
+        jax.random.key(0), self.qvalue.shape, dtype=self.qvalue.dtype
+    )
+    next_cotangent = bwd_fn(cotangent)
+
+    # Check that cotangents are correct
+    np.testing.assert_allclose(next_cotangent[0], cotangent)
+
+
+class FromHiQArrayTest(absltest.TestCase):
+  # Checks that from_qarray works
+
+  def setUp(self):
+    super().setUp()
+    orig_shape = (16, 16)
+    quant_axes = (0, 1)
+    group_sizes = (2, 2)
+    self.qarray = test_utils.create_random_qarray(
+        orig_shape, quant_axes, group_sizes, jax.random.key(0)
+    )
+
+  def test_from_qarray(self):
+    array = hq.from_hiqarray(self.qarray)
+    self.assertIsInstance(array, jax.Array)
+
+  def test_from_qarray_backward_with_zero_point(self):
+    # Set up simple examples for autograd testing
+    k1, k2, k3 = jax.random.split(jax.random.key(0), 3)
+    data = jax.random.normal(
+        k1, self.qarray.shape, dtype=self.qarray.metadata.dtype
+    )
+    scale = jnp.ones_like(self.qarray.scale)
+    zp = jnp.zeros_like(self.qarray.zero_point)
+    qarray = hq.to_hiqarray(data, scale, zp, self.qarray.metadata, key=k2)
+
+    arr = hq.from_hiqarray(qarray)
+    primals_out, bwd_fn = jax.vjp(hq.from_hiqarray, qarray)
+
+    # Test that primals are correct
+    np.testing.assert_allclose(arr, primals_out)
+
+    # Compute cotangents
+    cotangent = jax.random.normal(k3, arr.shape, dtype=arr.dtype)
+    next_cotangent = bwd_fn(cotangent)
+
+    # Check that cotangents are correct
+    np.testing.assert_allclose(next_cotangent[0], cotangent)
+
+
+class HiQArrayPermuteAxesTest(absltest.TestCase):
+
+  def setUp(self):
+    super().setUp()
+    orig_shape = (16, 16)
+    quant_axes = (0, 1)
+    group_sizes = (2, 2)
+    self.qarray = test_utils.create_random_qarray(
+        orig_shape, quant_axes, group_sizes, jax.random.key(0)
+    )
+
+  def test_permute_axes(self):
+    axes = (1, 0)
+    permuted_qarray = hq.permute_axes(self.qarray, axes)
+    self.assertIsInstance(permuted_qarray, hq.HiQArray)
+    np.testing.assert_allclose(
+        permuted_qarray.qvalue, jnp.permute_dims(self.qarray.qvalue, axes)
+    )
+    np.testing.assert_allclose(
+        permuted_qarray.scale, jnp.permute_dims(self.qarray.scale, axes)
+    )
+    np.testing.assert_allclose(
+        permuted_qarray.zero_point,
+        jnp.permute_dims(self.qarray.zero_point, axes),
+    )
+
+  def test_permute_axes_noop(self):
+    permuted_qarray = hq.permute_axes(self.qarray, (0, 1))
+    self.assertIsInstance(permuted_qarray, hq.HiQArray)
+    np.testing.assert_allclose(permuted_qarray.qvalue, self.qarray.qvalue)
+    np.testing.assert_allclose(permuted_qarray.scale, self.qarray.scale)
+    np.testing.assert_allclose(
+        permuted_qarray.zero_point, self.qarray.zero_point
+    )
+    self.assertEqual(permuted_qarray.metadata, self.qarray.metadata)
+
+  def test_permute_axes_involution(self):
+    permuted_qarray = hq.permute_axes(self.qarray, (1, 0))
+    permuted_qarray2 = hq.permute_axes(permuted_qarray, (1, 0))
+    np.testing.assert_allclose(permuted_qarray2.qvalue, self.qarray.qvalue)
+    np.testing.assert_allclose(permuted_qarray2.scale, self.qarray.scale)
+    np.testing.assert_allclose(
+        permuted_qarray2.zero_point, self.qarray.zero_point
+    )
+    self.assertEqual(permuted_qarray2.metadata, self.qarray.metadata)
+
+  def test_permute_axes_autograd(self):
+    cotangent = jax.random.normal(jax.random.key(0), self.qarray.shape)
+    primals_out, bwd_fn = jax.vjp(hq.permute_axes, self.qarray, (1, 0))
+    del primals_out  # Unused.
+    next_cotangent = bwd_fn(cotangent)
+    np.testing.assert_allclose(
+        next_cotangent[0], jnp.permute_dims(cotangent, (1, 0))
+    )
+
+
+class HiQArrayIntegrationTest(absltest.TestCase):
+
+  def test_to_from_qarray(self):
+    # Checks that round trip works for integer types and unit scale/zero point
+    orig_shape = (16, 16)
+    quant_info = {0: 2, 1: 2}
+    data = jax.random.uniform(
+        jax.random.key(0), orig_shape, minval=-128, maxval=128
+    ).astype(jnp.int32)
+    metadata = hqc.QuantizationMetadata.init(
+        orig_shape, quant_info, jnp.float32, jnp.int32
+    )
+    scale = jnp.ones(metadata.quant_shape, dtype=jnp.int32)
+    zero_point = jnp.zeros(metadata.quant_shape, dtype=jnp.int32)
+    qarray = hq.to_hiqarray(
+        data, scale, zero_point, metadata, key=jax.random.key(0)
+    )
+    array = hq.from_hiqarray(qarray)
+    qarray2 = hq.to_hiqarray(
+        array, scale, zero_point, metadata, key=jax.random.key(0)
+    )
+    np.testing.assert_allclose(qarray.qvalue, qarray2.qvalue)
+
+  def test_inspect(self):
+    orig_shape = (16, 16)
+    quant_info = {0: 2, 1: 2}
+    data = jax.random.uniform(
+        jax.random.key(0), orig_shape, minval=-128, maxval=128
+    ).astype(jnp.int32)
+    metadata = hqc.QuantizationMetadata.init(
+        orig_shape, quant_info, jnp.float32, jnp.int32
+    )
+    scale = jnp.ones(metadata.quant_shape, dtype=jnp.int32)
+    zero_point = jnp.zeros(metadata.quant_shape, dtype=jnp.int32)
+    fn = jax.make_jaxpr(hq.to_hiqarray, static_argnums=(3,))(
+        data, scale, zero_point, metadata, key=jax.random.key(0)
+    )
+    print(fn)
+
+
+if __name__ == "__main__":
+  absltest.main()