Feat (quant_tensor): convert from NamedTuple to Tensor subclass

[Giuseppe5]

Summary

This PR converts all QuantTensor types (IntQuantTensor, FloatQuantTensor,
GroupwiseIntQuantTensor, GroupwiseFloatQuantTensor) from NamedTuple subclasses
to proper torch.Tensor subclasses. The underlying tensor data now is the
dequantized value; quantization metadata (scale, zero_point, bit_width, etc.) are
stored as regular Python attributes exposed through property accessors.

---

Motivation

The previous NamedTuple-based design had several practical limitations:

• Tuple-decay in tracing: during ONNX export and torch.jit.trace, a QuantTensor
  could silently decay to a plain tuple, requiring a fragile get_quant_tensor_class
  reconstruction heuristic keyed on tuple length.
• No real Tensor semantics: isinstance(qt, torch.Tensor) was False, so any
  PyTorch API that dispatches on tensor type (e.g. torch.onnx.export argument
  handling, torch.compile graph tracing) had to be worked around with special cases.
• torch.compile / dynamo friction: metadata-only caching required storing None
  as the value placeholder, which is not compatible with tensor graph tracing.

Moving to a torch.Tensor subclass resolves all of the above cleanly:
QuantTensor instances are real tensors, participate naturally in PyTorch dispatch,
and no reconstruction heuristics are needed in tracing paths.

---

Changes

Core base class (src/brevitas/quant_tensor/base_quant_tensor.py)

• QuantTensor now subclasses torch.Tensor directly.
• __new__ creates the instance via value.as_subclass(cls), preserving grad_fn
  and requires_grad.
• value property returns a plain-Tensor view of self via
  Tensor.as_subclass(self, Tensor) to prevent infinite recursion in operations.
• Added _field_to_constructor_param class attribute for field-name → constructor-param
  mapping (used by groupwise types where e.g. scale_ maps to scale).
• Added _get_constructor_kwargs() helper that collects current metadata into a dict
  suitable for passing back to the constructor.
• set() is reimplemented (previously delegated to NamedTuple._replace()); it now
  builds a fresh instance via _get_constructor_kwargs() and handles the
  value/value_ distinction for groupwise types.
• detach(), contiguous(), to(), cuda(), cpu() are all reimplemented to
  iterate over _get_constructor_kwargs() and return a new instance of the same type
  with metadata tensors transformed accordingly.
• detach_() calls super().detach_() in addition to detaching each metadata tensor.
• Added __iadd__, __imul__, __isub__ (non-mutating; return a new value).
• Removed the four NamedTuple base classes: IntQuantTensorBase,
  FloatQuantTensorBase, GroupwiseFloatQuantTensorBase, GroupwisIntQuantTensorBase.

IntQuantTensor (src/brevitas/quant_tensor/int_quant_tensor.py)

• No longer inherits from IntQuantTensorBase.
• Declares _fields = ('scale', 'zero_point', 'bit_width', 'signed', 'training').
  Note: value is no longer a field — it is the tensor itself.
• __new__ performs value.as_subclass(cls); all metadata coercion moves to
  __init__, which stores _scale, _zero_point, _bit_width, signed_t,
  training_t.
• scale, zero_point, bit_width exposed as property getters/setters.

FloatQuantTensor (src/brevitas/quant_tensor/float_quant_tensor.py)

• No longer inherits from FloatQuantTensorBase.
• Declares _fields for all 10 metadata attributes (value excluded).
• Same __new__/__init__ split pattern as IntQuantTensor.
• All metadata attributes (scale, zero_point, exponent_bit_width,
  mantissa_bit_width, exponent_bias, inf_values, nan_values) exposed as
  property getters/setters.

GroupwiseIntQuantTensor (src/brevitas/quant_tensor/groupwise_int_quant_tensor.py)

• No longer inherits from GroupwisIntQuantTensorBase.
• Declares _fields (8 entries, using scale_/zero_point_ names for the grouped
  tensors) and _field_to_constructor_param = {'scale_': 'scale', 'zero_point_': 'zero_point'}.
• Sets _is_groupwise = True so that set()/detach()/etc. in the base class use
  _value_ instead of value.
• The raw grouped value is stored as self._value_; self (as a tensor) holds the
  dequantized/expanded representation used by .value.
• group_size, group_dim, bit_width, dequant_shape exposed as properties.
• Fixed __neg__ and __abs__ to reference self.zero_point_ / self.scale_
  instead of the removed self.zero_point / self.scale.
• Removed incorrect saturating property that was accidentally inherited.

GroupwiseFloatQuantTensor (src/brevitas/quant_tensor/groupwise_float_quant_tensor.py)

• Same structural changes as GroupwiseIntQuantTensor.
• _fields covers 13 metadata attributes.
• All float metadata (exponent_bit_width, mantissa_bit_width, exponent_bias,
  inf_values, nan_values, dequant_shape) exposed as property getters/setters.
• Fixed __neg__ and __abs__ to use self.zero_point_ / self.scale_.

__init__.py (src/brevitas/quant_tensor/__init__.py)

• Changed from from .base_quant_tensor import * (which exported the now-deleted
  NamedTuple base classes) to explicit named imports of _unpack_quant_tensor and
  QuantTensor.

NN mixin (src/brevitas/nn/mixin/base.py)

• Removed get_quant_tensor_class() — the tuple-length heuristic for reconstructing
  a QuantTensor from a decayed tuple is no longer needed.
• Removed the tracing-state tuple-decay workaround in unpack_input(). Since
  QuantTensor is now a real Tensor, it does not decay to a plain tuple during
  tracing.
• Removed unused imports of FloatQuantTensor, GroupwiseFloatQuantTensor,
  GroupwiseIntQuantTensor from this module.

ONNX export (src/brevitas/export/onnx/manager.py)

• Simplified the isinstance(args, tuple) guard: previously it was
  isinstance(args, tuple) and not isinstance(args, QuantTensor) because
  QuantTensor was a NamedTuple (hence a tuple). This exclusion is no longer
  needed.

Proxy (src/brevitas/proxy/parameter_quant.py)

• Updated the cached-activation fallback check from
  isinstance(quant_input, Tensor) to not isinstance(quant_input, IntQuantTensor),
  since IntQuantTensor is now itself a Tensor subclass and the old check would
  always be True.

Caching utilities (src/brevitas/utils/quant_utils.py)

• _CachedIO, _CachedIOFloat, _CachedIOGroupwiseFloat, _CachedIOGroupwiseInt:
  metadata-only caching now stores torch.empty(0) as the value placeholder instead
  of None. A None value is not valid inside a torch.Tensor subclass and breaks
  torch.compile graph tracing.

Notebook (notebooks/minifloat_mx_tutorial.ipynb)

• Minor output cell refresh (3 lines updated).

---

Behavioral / Migration Notes

Code that worked with the old NamedTuple-based QuantTensor may need updates in
the following areas:

Old behaviour	New behaviour
isinstance(qt, tuple) → True	isinstance(qt, tuple) → False
isinstance(qt, torch.Tensor) → False	isinstance(qt, torch.Tensor) → True
len(qt._fields) == 6 for IntQuantTensor	len(qt._fields) == 5 (value is no longer a field)
qt._replace(scale=s) (NamedTuple API)	qt.set(scale=s) (already the public API)
qt.tensor (alias for .value)	Removed; use .value directly
qt.set(value=None) for metadata-only caching	qt.set(value=torch.empty(0))
qt[0], qt[1], … index access	Not supported; use named properties instead
QuantTensor decays to tuple during tracing	Does not decay; tracing treats it as a Tensor

The public .value, .scale, .zero_point, .bit_width (and float equivalents)
property API is unchanged.

---

Testing

• tests/brevitas/quant_tensor/test_quant_tensor.py — updated assertions:
  isinstance(qt, torch.Tensor) (was tuple); len(qt._fields) == 5 (was 6).

• tests/brevitas/quant_tensor/test_quant_tensor_subclass.py (new, 513 lines) —
  comprehensive test suite covering the new subclass behaviour:

  • TestIsInstance — IntQuantTensor, FloatQuantTensor, GroupwiseFloatQuantTensor
    are all isinstance of torch.Tensor and QuantTensor.
  • TestValueProperty — .value returns a plain torch.Tensor (not a subclass);
    .value and self share data; grad_fn and requires_grad are preserved.
  • TestSetMethod — set() replaces fields correctly, preserves type, no-args copy.
  • TestInPlaceOperators — +=, *=, -= return valid results and preserve type.
  • TestTorchFunctionFallback — unhandled functions fall back to plain tensor output
    without recursion; handled functions (relu, max_pool2d) route correctly.
  • TestTensorOps — detach(), contiguous(), to(dtype) preserve type and metadata.
  • TestUnpackQuantTensor — plain tensor passthrough; tuple/dict/nested unpacking.
  • TestRightHandOperators — __radd__, __rmul__ with plain-tensor LHS.
  • TestFields — _fields contents for Int, Float, GroupwiseFloat variants.
  • TestShapeSizeDim — .shape, .size(), .dim() delegate correctly.
  • TestConstructionFromNonTensor — construction from Python float or list.
  • TestDynamoExportCacheClass — _CachedIO metadata-only path works end-to-end.

[Giuseppe5]

We should remove quant_tensor.tensor since it was there for backwards compatibility but it was deprecated.