test propagation using simple buffer casting transform test

Author: Ali Alshaarawy

thunder/tests/test_transforms.py

@@ -647,21 +647,22 @@ def test_disable_params_and_buffer_check():
 
     assert len(check_bsyms) == 1  # We only have the check for input.
 
-def test_buffer_quantization():
+
+def test_buffer_dtype_casting():
     import torch.nn as nn
     import itertools
 
     from typing import Any, Optional, Tuple, Union, List
 
     class CastBuffers(thunder.core.transform_common.Transform):
         def __init__(self):
-            self.quant_states = {}
+            self.cast_states = {}
 
         def transform_module(self, model: thunder.ThunderModule):
             self.thunder_module = model
             for n, b in model._model.named_buffers():
                 qb = b.to(torch.bfloat16)
-                self.quant_states[n] = {
+                self.cast_states[n] = {
                     "dtype": b.dtype,
                     "shape": tuple(b.shape),
                     "qb.dtype": qb.dtype,
@@ -679,13 +680,13 @@ def transform_traces_pre_prologue(
 
             prologue_proxy_map = {
                 get_param_bsym.output.name: dict(
-                    shape=self.quant_states[model_weight_name]["qb.shape"],
+                    shape=self.cast_states[model_weight_name]["qb.shape"],
                     dtype=thunder.dtypes.to_dtype(
-                        self.quant_states[model_weight_name]["qb.dtype"]
+                        self.cast_states[model_weight_name]["qb.dtype"]
                     ),
                 )
                 for model_weight_name, (check_bsym, get_param_bsym) in checks.items()
-                if model_weight_name in self.quant_states
+                if model_weight_name in self.cast_states
             }
 
             # here we switch the prologue_trace to a copy with new metadata
@@ -696,7 +697,7 @@ def transform_traces_pre_prologue(
             )
 
             checks = thunder.transforms.utils.get_checks(prologue_trace)
-            for n, qs in self.quant_states.items():
+            for n, qs in self.cast_states.items():
                 check, get_param = checks[n]
                 # check has args: tensor, shape, device, dtype, requires_grad
                 proxy, _, device, _, requires_grad = check.args
@@ -823,4 +824,4 @@ def check_dtypes(bsym):
     for tr in thunder.last_traces(cast_jit):
         if str(tr.get_provenance()) == "# Constructed by Dtype Convert":
             for bsym in tr.bound_symbols:
-                check_dtypes(bsym)
+                check_dtypes(bsym)