An extension of the torch.Tensor class to support our inference focused features. One important thing to note here is that an InferenceParam can be used a torch.Tensor, but outputs of torch.Tensor operations will not be InferenceParams.
| 24 | |
| 25 | |
| 26 | class InferenceParameter(torch.Tensor): |
| 27 | """ |
| 28 | An extension of the torch.Tensor class to support our inference focused features. One important |
| 29 | thing to note here is that an InferenceParam can be used a torch.Tensor, but outputs of |
| 30 | torch.Tensor operations will not be InferenceParams. |
| 31 | """ |
| 32 | |
| 33 | @staticmethod |
| 34 | def __new__(cls, tensor, *args, **kwargs): |
| 35 | new_tensor = super().__new__(cls, tensor, *args, **kwargs) |
| 36 | if hasattr(tensor, "_aux_attrs"): |
| 37 | setattr(new_tensor, "_aux_attrs", tensor.aux_attrs) |
| 38 | return new_tensor |
| 39 | |
| 40 | def to(self, *args, **kwargs): |
| 41 | new_tensor = super().to(*args, **kwargs) |
| 42 | if hasattr(self, "_aux_attrs"): |
| 43 | setattr(new_tensor, "_aux_attrs", self.aux_attrs) |
| 44 | try: |
| 45 | _ = torch.device(args[0]) |
| 46 | for name, attr in new_tensor.aux_attrs.items(): |
| 47 | new_attr = attr.to(*args, **kwargs) |
| 48 | setattr(new_tensor, name, new_attr) |
| 49 | new_tensor.aux_attrs[name] = new_attr |
| 50 | except (RuntimeError, TypeError, IndexError): |
| 51 | pass |
| 52 | |
| 53 | return new_tensor |
| 54 | |
| 55 | @classmethod |
| 56 | def initialize(cls, core_param: torch.Tensor, **kwargs) -> 'InferenceParameter': |
| 57 | """ |
| 58 | Create the inference parameter. |
| 59 | """ |
| 60 | param = InferenceParameter(core_param) |
| 61 | setattr(param, "_aux_attrs", kwargs) |
| 62 | |
| 63 | for attr_name, attr in kwargs.items(): |
| 64 | if hasattr(param, attr_name): |
| 65 | raise ValueError(f"Attribute {attr_name} already exists on param.") |
| 66 | |
| 67 | if not isinstance(attr, torch.Tensor): |
| 68 | raise ValueError(f"Attribute {attr_name} must be a tensor.") |
| 69 | |
| 70 | setattr(param, attr_name, attr) |
| 71 | |
| 72 | return param |
| 73 | |
| 74 | @classmethod |
| 75 | def initialize_raw(self, **kwargs) -> 'InferenceParameter': |
| 76 | """ |
| 77 | All kwargs must be torch.Tensors and must include the core parameter. |
| 78 | """ |
| 79 | if CORE_PARAM not in kwargs: |
| 80 | raise ValueError(f"Must provide core parameter, with key {CORE_PARAM}.") |
| 81 | |
| 82 | return InferenceParameter.initialize(kwargs[CORE_PARAM], **kwargs) |
| 83 |