diff --git a/bofire/benchmarks/benchmark.py b/bofire/benchmarks/benchmark.py
index ff1f24835..dac00df2d 100644
--- a/bofire/benchmarks/benchmark.py
+++ b/bofire/benchmarks/benchmark.py
@@ -135,31 +135,28 @@ def __init__(
     ):
         super().__init__(**kwargs)
         self._benchmark = benchmark
+        benchmark_inputs = self._benchmark.domain.inputs.get()
         assert n_filler_features >= 1, "n_filler_features must be >= 1."
         assert len(benchmark.domain.constraints) == 0, "Constraints not supported yet."
-        assert len(benchmark.domain.inputs.get(ContinuousInput)) == len(
-            benchmark.domain.inputs
-        ), "Only continuous inputs supported yet."
+        if not Inputs.is_continuous(benchmark_inputs):
+            raise ValueError("Only continuous inputs supported yet.")
         self.n_filler_features = n_filler_features
         self.n_features_per_original_feature = n_features_per_original_feature
 
         features = []
         constraints = []
-        for j, feat in enumerate(self._benchmark.domain.inputs.get()):
+        for j, feat in enumerate(benchmark_inputs):
             features += [
                 ContinuousInput(
                     key=f"{feat.key}_{i}",
-                    bounds=(0, 1 / len(self._benchmark.domain.inputs)),
+                    bounds=(0, 1 / len(benchmark_inputs)),
                 )
                 if self.n_features_per_original_feature == 1
                 else ContinuousDescriptorInput(
                     key=f"{feat.key}_{i}",
                     bounds=(0, 1 / len(self._benchmark.domain.inputs)),
                     descriptors=self._benchmark.domain.inputs.get_keys(),
-                    values=[
-                        1 if k == j else 0
-                        for k in range(len(self._benchmark.domain.inputs))
-                    ],
+                    values=[1 if k == j else 0 for k in range(len(benchmark_inputs))],
                 )
                 for i in range(self.n_features_per_original_feature)
             ]
@@ -171,7 +168,7 @@ def __init__(
                             for i in range(self.n_features_per_original_feature)
                         ],
                         coefficients=[1.0] * self.n_features_per_original_feature,
-                        rhs=1 / len(self._benchmark.domain.inputs),
+                        rhs=1 / len(benchmark_inputs),
                     )
                 )
 
@@ -207,18 +204,13 @@ def __init__(
             constraints=Constraints(constraints=constraints),
             outputs=self._benchmark.domain.outputs,
         )
-        self._mins = np.array(
-            [feat.bounds[0] for feat in self._benchmark.domain.inputs.get()]
-        )
+
+        self._mins = np.array([feat.bounds[0] for feat in benchmark_inputs])
         self._scales = np.array(
-            [
-                feat.bounds[1] - feat.bounds[0]
-                for feat in self._benchmark.domain.inputs.get()
-            ]
+            [feat.bounds[1] - feat.bounds[0] for feat in benchmark_inputs]
         )
         self._scales_new = np.array(
-            [1 / len(self._benchmark.domain.inputs.get_keys())]
-            * len(self._benchmark.domain.inputs.get_keys())
+            [1 / len(benchmark_inputs.get_keys())] * len(benchmark_inputs.get_keys())
         )
 
     def _transform(self, X: pd.DataFrame) -> pd.DataFrame:
diff --git a/bofire/data_models/domain/constraints.py b/bofire/data_models/domain/constraints.py
index a7bc1fd77..7e8365d27 100644
--- a/bofire/data_models/domain/constraints.py
+++ b/bofire/data_models/domain/constraints.py
@@ -1,38 +1,44 @@
 import collections.abc
 from collections.abc import Iterator, Sequence
 from itertools import chain
-from typing import Generic, List, Literal, Optional, Type, TypeVar, Union
+from typing import Generic, Literal, Type, Union, overload
 
 import pandas as pd
 from pydantic import Field
+from typing_extensions import Self, TypeVar
 
 from bofire.data_models.base import BaseModel
 from bofire.data_models.constraints.api import AnyConstraint, Constraint
 from bofire.data_models.filters import filter_by_class
 
 
-C = TypeVar("C", bound=Union[AnyConstraint, Constraint])
-CIncludes = TypeVar("CIncludes", bound=Union[AnyConstraint, Constraint])
-CExcludes = TypeVar("CExcludes", bound=Union[AnyConstraint, Constraint])
+ConstraintT = TypeVar(
+    "ConstraintT", bound=AnyConstraint | Constraint, default=AnyConstraint | Constraint
+)
+ConstraintGetT = TypeVar(
+    "ConstraintGetT",
+    bound=AnyConstraint | Constraint,
+    default=AnyConstraint | Constraint,
+)
 
 
-class Constraints(BaseModel, Generic[C]):
+class Constraints(BaseModel, Generic[ConstraintT]):
     type: Literal["Constraints"] = "Constraints"
-    constraints: Sequence[C] = Field(default_factory=list)
+    constraints: Sequence[ConstraintT] = Field(default_factory=list)
 
-    def __iter__(self) -> Iterator[C]:
+    def __iter__(self) -> Iterator[ConstraintT]:
         return iter(self.constraints)
 
     def __len__(self):
         return len(self.constraints)
 
-    def __getitem__(self, i) -> C:
+    def __getitem__(self, i) -> ConstraintT:
         return self.constraints[i]
 
     def __add__(
         self,
-        other: Union[Sequence[CIncludes], "Constraints[CIncludes]"],
-    ) -> "Constraints[Union[C, CIncludes]]":
+        other: Union[Sequence[ConstraintGetT], "Constraints[ConstraintGetT]"],
+    ) -> "Constraints[Union[ConstraintT, ConstraintGetT]]":
         if isinstance(other, collections.abc.Sequence):
             other_constraints = other
         else:
@@ -87,14 +93,28 @@ def is_fulfilled(self, experiments: pd.DataFrame, tol: float = 1e-6) -> pd.Serie
             .all(axis=1)
         )
 
+    @overload
     def get(
         self,
-        includes: Union[
-            Type[CIncludes], Sequence[Type[CIncludes]]
-        ] = Constraint,  # ty: ignore[invalid-parameter-default]
-        excludes: Optional[Union[Type[CExcludes], List[Type[CExcludes]]]] = None,
+        includes: Type[ConstraintGetT] | Sequence[Type[ConstraintGetT]],
+        excludes: None = None,
         exact: bool = False,
-    ) -> "Constraints[CIncludes]":
+    ) -> "Constraints[ConstraintGetT]": ...
+
+    @overload
+    def get(
+        self,
+        includes: Type[ConstraintGetT] | Sequence[Type[ConstraintGetT]] = Constraint,
+        excludes: Type[AnyConstraint] | Sequence[Type[AnyConstraint]] | None = None,
+        exact: bool = False,
+    ) -> Self: ...
+
+    def get(
+        self,
+        includes: Type[ConstraintGetT] | Sequence[Type[ConstraintGetT]] = Constraint,
+        excludes: Type[AnyConstraint] | Sequence[Type[AnyConstraint]] | None = None,
+        exact: bool = False,
+    ) -> "Constraints[ConstraintGetT]":
         """Get constraints of the domain
 
         Args:
diff --git a/bofire/data_models/domain/features.py b/bofire/data_models/domain/features.py
index 6a350bbd1..30ddd11ab 100644
--- a/bofire/data_models/domain/features.py
+++ b/bofire/data_models/domain/features.py
@@ -15,16 +15,16 @@
     Optional,
     Tuple,
     Type,
-    TypeVar,
     Union,
     cast,
+    overload,
 )
 
 import numpy as np
 import pandas as pd
 from pydantic import Field, field_validator, validate_call
 from scipy.stats.qmc import LatinHypercube, Sobol
-from typing_extensions import Self
+from typing_extensions import Self, TypeGuard, TypeVar
 
 from bofire.data_models.base import BaseModel
 from bofire.data_models.enum import CategoricalEncodingEnum, SamplingMethodEnum
@@ -56,11 +56,20 @@
 from bofire.data_models.types import InputTransformSpecs
 
 
-F = TypeVar("F", bound=AnyFeature)
-FeatureSequence = Sequence[F]
+FeatureT = TypeVar("FeatureT", bound=AnyFeature | Feature, default=AnyFeature | Feature)
+FeatureGetT = TypeVar(
+    "FeatureGetT", bound=AnyFeature | Feature, default=AnyFeature | Feature
+)
+InputT = TypeVar("InputT", bound=AnyInput | Input, default=AnyInput | Input)
+InputGetT = TypeVar("InputGetT", bound=AnyInput | Input, default=AnyInput | Input)
+EngineeredFeatureT = TypeVar(
+    "EngineeredFeatureT", bound=AnyEngineeredFeature, default=AnyEngineeredFeature
+)
+OutputT = TypeVar("OutputT", bound=AnyOutput | Output, default=AnyOutput | Output)
+OutputGetT = TypeVar("OutputGetT", bound=AnyOutput | Output, default=AnyOutput | Output)
 
 
-class _BaseFeatures(BaseModel, Generic[F]):
+class _BaseFeatures(BaseModel, Generic[FeatureT]):
     """Container of features, both input and output features are allowed.
 
     Attributes:
@@ -69,29 +78,29 @@ class _BaseFeatures(BaseModel, Generic[F]):
     """
 
     type: Literal["Features"] = "Features"
-    features: FeatureSequence = Field(default_factory=list)
+    features: Sequence[FeatureT] = Field(default_factory=list)
 
     @field_validator("features")
     @classmethod
     def validate_unique_feature_keys(
         cls: type[_BaseFeatures],
-        features: FeatureSequence,
-    ) -> FeatureSequence:
+        features: Sequence[FeatureT],
+    ) -> Sequence[FeatureT]:
         keys = [feat.key for feat in features]
         if len(keys) != len(set(keys)):
             raise ValueError("Feature keys are not unique.")
         return features
 
-    def __iter__(self) -> Iterator[F]:
+    def __iter__(self) -> Iterator[FeatureT]:
         return iter(self.features)
 
     def __len__(self):
         return len(self.features)
 
-    def __getitem__(self, i):
+    def __getitem__(self, i) -> FeatureT:
         return self.features[i]
 
-    def __add__(self, other: Union[Sequence[AnyFeature], _BaseFeatures]):
+    def __add__(self, other: Union[Sequence[AnyFeature], _BaseFeatures[FeatureGetT]]):
         if isinstance(other, Features):
             other_feature_seq = other.features
         else:
@@ -126,7 +135,7 @@ def is_engineeredfeats(feats):
             )
         return Features(features=new_feature_seq)
 
-    def get_by_key(self, key: str, use_regex: bool = False) -> F:
+    def get_by_key(self, key: str, use_regex: bool = False) -> FeatureT:
         """Get a feature by its key.
 
         First, the method tries to find the feature by its key. If no feature is
@@ -174,10 +183,8 @@ def get_by_keys(self, keys: Sequence[str], include: bool = True) -> Self:
 
     def get(
         self,
-        includes: Union[
-            Type, List[Type], None
-        ] = AnyFeature,  # ty: ignore[invalid-parameter-default]
-        excludes: Union[Type, List[Type], None] = None,
+        includes: Type[FeatureGetT] | Sequence[Type[FeatureGetT]] = Feature,
+        excludes: Type[AnyFeature] | Sequence[Type[AnyFeature]] | None = None,
         exact: bool = False,
     ) -> Self:
         """Get features of this container and filter via includes and excludes.
@@ -208,17 +215,15 @@ def get(
 
     def get_keys(
         self,
-        includes: Union[
-            Type, List[Type], None
-        ] = AnyFeature,  # ty: ignore[invalid-parameter-default]
-        excludes: Union[Type, List[Type], None] = None,
+        includes: Type[FeatureGetT] | Sequence[Type[FeatureGetT]] = Feature,
+        excludes: Type[AnyFeature] | Sequence[Type[AnyFeature]] | None = None,
         exact: bool = False,
-    ) -> List[str]:
+    ) -> list[str]:
         """Get feature-keys of this container and filter via includes and excludes.
 
         Args:
             includes: All features in this container that are instances of an
-                include are returned. If None, the include filter is not active.
+                include are returned.
             excludes: All features in this container that are not instances of
                 an exclude are returned. If None, the exclude filter is not active.
             exact: Boolean to distinguish if only the exact class listed in
@@ -251,11 +256,11 @@ def get_reps_df(self) -> pd.DataFrame:
         return df
 
 
-class Features(_BaseFeatures[AnyFeature]):
+class Features(_BaseFeatures[FeatureT]):
     pass
 
 
-class EngineeredFeatures(_BaseFeatures[AnyEngineeredFeature]):
+class EngineeredFeatures(_BaseFeatures[EngineeredFeatureT]):
     """Container of engineered (input) features, only engineered features
     are allowed.
 
@@ -331,7 +336,7 @@ def n_transformed_inputs(self) -> int:
         return sum(feat.n_transformed_inputs for feat in self.get())
 
 
-class Inputs(_BaseFeatures[AnyInput]):
+class Inputs(_BaseFeatures[InputT]):
     """Container of input features, only input features are allowed.
 
     Attributes:
@@ -343,7 +348,7 @@ class Inputs(_BaseFeatures[AnyInput]):
 
     @field_validator("features")
     @classmethod
-    def validate_only_one_task_input(cls, features: Sequence[AnyInput]):
+    def validate_only_one_task_input(cls, features: Sequence[InputT]):
         filtered = filter_by_class(
             features,
             includes=TaskInput,
@@ -412,7 +417,7 @@ def sample(
                 pd.concat(
                     [
                         feat.sample(n, seed=int(rng.integers(1, 1000000)))
-                        for feat in self.get(Input)
+                        for feat in self.get()
                     ],
                     axis=1,
                 ),
@@ -501,10 +506,8 @@ def validate_experiments(
 
     def get_number_of_categorical_combinations(
         self,
-        include: Union[Type, List[Type]] = Input,
-        exclude: Union[
-            Type, List[Type]
-        ] = None,  # ty: ignore[invalid-parameter-default]
+        include: Type[InputGetT] | list[Type[InputGetT]] = Input,
+        exclude: Type[AnyInput] | list[Type[AnyInput]] | None = None,
     ) -> int:
         """Get the total number of unique categorical combinations.
 
@@ -549,10 +552,8 @@ def get_number_of_categorical_combinations(
 
     def get_categorical_combinations(
         self,
-        include: Union[Type, List[Type]] = Input,
-        exclude: Union[
-            Type, List[Type]
-        ] = None,  # ty: ignore[invalid-parameter-default]
+        include: Type[InputGetT] | list[Type[InputGetT]] = Input,
+        exclude: Type[AnyInput] | List[Type[AnyInput]] | None = None,
     ) -> list[tuple[tuple[str, float] | tuple[str, str], ...]]:
         """Get a list of tuples pairing the feature keys with a list of valid categories
 
@@ -629,7 +630,7 @@ def _get_transform_info(
         features2idx = {}
         features2names = {}
         counter = 0
-        for _, feat in enumerate(self.get()):
+        for feat in self.get():
             if feat.key not in specs.keys():
                 features2idx[feat.key] = (counter,)
                 features2names[feat.key] = (feat.key,)
@@ -911,8 +912,37 @@ def is_fulfilled(self, experiments: pd.DataFrame) -> pd.Series:
             .all(axis=1)
         )
 
+    @staticmethod
+    def is_continuous(inputs: Inputs) -> TypeGuard[Inputs[ContinuousInput]]:
+        return len(inputs.get(ContinuousInput)) == len(inputs)
+
+    @overload
+    def get(
+        self,
+        includes: Type[InputGetT] | Sequence[Type[InputGetT]] = Input,
+        excludes: None = None,
+        exact: bool = False,
+    ) -> Inputs[InputGetT]: ...
+
+    @overload
+    def get(
+        self,
+        includes: Type[InputGetT] | Sequence[Type[InputGetT]],
+        excludes: Type[AnyInput] | Sequence[Type[AnyInput]] | None,
+        exact: bool = False,
+    ) -> Self: ...
+
+    def get(
+        self,
+        includes: Type[InputGetT] | Sequence[Type[InputGetT]] = Input,
+        excludes: Type[AnyInput] | Sequence[Type[AnyInput]] | None = None,
+        exact: bool = False,
+    ) -> Self:
+        # repeat the function here as implementation must be below overloads
+        return super().get(includes, excludes, exact)
+
 
-class Outputs(_BaseFeatures[AnyOutput]):
+class Outputs(_BaseFeatures[OutputT]):
     """Container of output features, only output features are allowed.
 
     Attributes:
@@ -1110,7 +1140,6 @@ def validate_candidates(self, candidates: pd.DataFrame) -> pd.DataFrame:
                     [f"{key}_pred", f"{key}_sd"]
                     for key in self.get_keys_by_objective(
                         excludes=Objective,
-                        includes=None,
                     )
                 ],
             ),
@@ -1215,3 +1244,28 @@ def preprocess_experiments_any_valid_output(
             ),
         )
         return clean_exp
+
+    @overload
+    def get(
+        self,
+        includes: Type[OutputGetT] | Sequence[Type[OutputGetT]],
+        excludes: None = None,
+        exact: bool = False,
+    ) -> Outputs[OutputGetT]: ...
+
+    @overload
+    def get(
+        self,
+        includes: Type[OutputGetT] | Sequence[Type[OutputGetT]],
+        excludes: Type[AnyOutput] | Sequence[Type[AnyOutput]] | None,
+        exact: bool = False,
+    ) -> Self: ...
+
+    def get(
+        self,
+        includes: Type[OutputGetT] | Sequence[Type[OutputGetT]] = Output,
+        excludes: Type[AnyOutput] | Sequence[Type[AnyOutput]] | None = None,
+        exact: bool = False,
+    ) -> Self:
+        # repeat the function here as implementation must be below overloads
+        return super().get(includes, excludes, exact)
diff --git a/bofire/data_models/features/feature.py b/bofire/data_models/features/feature.py
index f66d0c533..ab21cbf54 100644
--- a/bofire/data_models/features/feature.py
+++ b/bofire/data_models/features/feature.py
@@ -152,6 +152,8 @@ class Output(Feature):
 
     """
 
+    objective: Any
+
     @abstractmethod
     def __call__(self, values: pd.Series) -> pd.Series:
         pass
diff --git a/bofire/data_models/surrogates/deterministic.py b/bofire/data_models/surrogates/deterministic.py
index 31cee404e..5e4b5e595 100644
--- a/bofire/data_models/surrogates/deterministic.py
+++ b/bofire/data_models/surrogates/deterministic.py
@@ -2,7 +2,7 @@
 
 from pydantic import Field, field_validator, model_validator
 
-from bofire.data_models.domain.api import EngineeredFeatures
+from bofire.data_models.domain.api import EngineeredFeatures, Inputs
 from bofire.data_models.features.api import (
     AnyOutput,
     CategoricalInput,
@@ -27,6 +27,7 @@ class CategoricalDeterministicSurrogate(BotorchSurrogate):
     type: Literal["CategoricalDeterministicSurrogate"] = (
         "CategoricalDeterministicSurrogate"
     )
+    inputs: Inputs[CategoricalInput]
     mapping: Annotated[Dict[str, float], Field(min_length=2)]
 
     @model_validator(mode="after")
diff --git a/bofire/strategies/predictives/acqf_optimization.py b/bofire/strategies/predictives/acqf_optimization.py
index 0a5938397..0fde178ff 100644
--- a/bofire/strategies/predictives/acqf_optimization.py
+++ b/bofire/strategies/predictives/acqf_optimization.py
@@ -29,7 +29,6 @@
     CategoricalInput,
     ContinuousInput,
     DiscreteInput,
-    Input,
 )
 from bofire.data_models.strategies.api import (
     AcquisitionOptimizer as AcquisitionOptimizerDataModel,
@@ -207,8 +206,7 @@ def get_fixed_features(
             domain
         )
 
-        for _, feat in enumerate(domain.inputs.get(Input)):
-            assert isinstance(feat, Input)
+        for _, feat in enumerate(domain.inputs.get()):
             if feat.fixed_value() is not None:
                 fixed_values = feat.fixed_value(
                     transform_type=input_preprocessing_specs.get(feat.key),
diff --git a/bofire/strategies/predictives/multi_fidelity.py b/bofire/strategies/predictives/multi_fidelity.py
index 18af1086b..4c46cd267 100644
--- a/bofire/strategies/predictives/multi_fidelity.py
+++ b/bofire/strategies/predictives/multi_fidelity.py
@@ -118,9 +118,9 @@ def _verify_all_fidelities_observed(self) -> None:
         assert self.experiments is not None
         observed_fidelities = set(self.experiments[self.task_feature_key].unique())
         allowed_fidelities = set(
-            self.domain.inputs.get_by_key(
-                self.task_feature_key
-            ).get_allowed_categories()
+            self.domain.inputs.get(TaskInput)
+            .get_by_key(self.task_feature_key)
+            .get_allowed_categories()
         )
         missing_fidelities = allowed_fidelities - observed_fidelities
         if missing_fidelities:
diff --git a/bofire/strategies/shortest_path.py b/bofire/strategies/shortest_path.py
index b0a4afef1..6adae3924 100644
--- a/bofire/strategies/shortest_path.py
+++ b/bofire/strategies/shortest_path.py
@@ -47,7 +47,7 @@ def continuous_inputs(self) -> Inputs:
 
     def get_linear_constraints(
         self,
-        constraints: Constraints,
+        constraints: Constraints[LinearConstraint],
     ) -> Tuple[np.ndarray, np.ndarray]:
         """Returns the linear constraints in the form of matrices A and b, where Ax = b for
             equality constraints and Ax <= b for inequality constraints.
diff --git a/bofire/strategies/stepwise/stepwise.py b/bofire/strategies/stepwise/stepwise.py
index c5dc850d1..0fbee1c7a 100644
--- a/bofire/strategies/stepwise/stepwise.py
+++ b/bofire/strategies/stepwise/stepwise.py
@@ -51,7 +51,7 @@ def get_step(self) -> Tuple[Strategy, Optional[Transform]]:
                 ]  # ty: ignore[invalid-return-type]
         raise ValueError("No condition could be satisfied.")
 
-    def _ask(self, candidate_count: Optional[PositiveInt]) -> pd.DataFrame:
+    def _ask(self, candidate_count: Optional[PositiveInt] = None) -> pd.DataFrame:
         strategy, transform = self.get_step()
 
         candidate_count = candidate_count or 1
diff --git a/bofire/surrogates/deterministic.py b/bofire/surrogates/deterministic.py
index 3871534c2..9b7ccb1ea 100644
--- a/bofire/surrogates/deterministic.py
+++ b/bofire/surrogates/deterministic.py
@@ -62,7 +62,7 @@ def __init__(
         self.model = AffineDeterministicModel(
             b=0.0,
             a=torch.tensor(
-                [data_model.mapping[key] for key in self.inputs[0].categories],
+                [data_model.mapping[key] for key in data_model.inputs[0].categories],
             )
             .to(**tkwargs)
             .unsqueeze(-1),
diff --git a/bofire/surrogates/trainable.py b/bofire/surrogates/trainable.py
index a66e7984b..d03b24fca 100644
--- a/bofire/surrogates/trainable.py
+++ b/bofire/surrogates/trainable.py
@@ -221,29 +221,19 @@ def cross_validate(
             y_train_pred = self.predict(X_train)
 
             # Convert to categorical if applicable
-            if isinstance(
-                self.outputs.get_by_key(key).objective,
-                ConstrainedCategoricalObjective,
-            ):
+            objective = self.outputs.get_by_key(key)
+            if isinstance(objective, ConstrainedCategoricalObjective):
                 y_test_pred[f"{key}_pred"] = y_test_pred[f"{key}_pred"].map(
-                    self.outputs.get_by_key(
-                        key
-                    ).objective.to_dict_label(),  # ty: ignore[possibly-missing-attribute]
+                    objective.to_dict_label(),
                 )
                 y_train_pred[f"{key}_pred"] = y_train_pred[f"{key}_pred"].map(
-                    self.outputs.get_by_key(
-                        key
-                    ).objective.to_dict_label(),  # ty: ignore[possibly-missing-attribute]
+                    objective.to_dict_label(),
                 )
                 y_test[key] = y_test[key].map(
-                    self.outputs.get_by_key(
-                        key
-                    ).objective.to_dict_label(),  # ty: ignore[possibly-missing-attribute]
+                    objective.to_dict_label(),
                 )
                 y_train[key] = y_train[key].map(
-                    self.outputs.get_by_key(
-                        key
-                    ).objective.to_dict_label(),  # ty: ignore[possibly-missing-attribute]
+                    objective.to_dict_label(),
                 )
 
             # now store the results
diff --git a/bofire/utils/naming_conventions.py b/bofire/utils/naming_conventions.py
index 645a692fe..29dd6ba46 100644
--- a/bofire/utils/naming_conventions.py
+++ b/bofire/utils/naming_conventions.py
@@ -17,13 +17,9 @@ def get_column_names(outputs: Outputs) -> Tuple[List[str], List[str]]:
 
     """
     pred_cols, sd_cols = [], []
-    for featkey in outputs.get_keys(CategoricalOutput):
-        pred_cols = pred_cols + [
-            f"{featkey}_{cat}_prob" for cat in outputs.get_by_key(featkey).categories
-        ]
-        sd_cols = sd_cols + [
-            f"{featkey}_{cat}_sd" for cat in outputs.get_by_key(featkey).categories
-        ]
+    for feat in outputs.get(CategoricalOutput):
+        pred_cols = pred_cols + [f"{feat.key}_{cat}_prob" for cat in feat.categories]
+        sd_cols = sd_cols + [f"{feat.key}_{cat}_sd" for cat in feat.categories]
     for featkey in outputs.get_keys(ContinuousOutput):
         pred_cols = pred_cols + [f"{featkey}_pred"]
         sd_cols = sd_cols + [f"{featkey}_sd"]
diff --git a/bofire/utils/reduce.py b/bofire/utils/reduce.py
index 483da8654..a8db45c73 100644
--- a/bofire/utils/reduce.py
+++ b/bofire/utils/reduce.py
@@ -95,9 +95,7 @@ def reduce_domain(domain: Domain) -> Tuple[Domain, AffineTransform]:
     )
 
     # only consider continuous inputs
-    continuous_inputs = [
-        cast(ContinuousInput, f) for f in domain.inputs.get(ContinuousInput)
-    ]
+    continuous_inputs = list(domain.inputs.get(ContinuousInput))
     other_inputs = domain.inputs.get(Input, excludes=[ContinuousInput])
 
     # assemble Matrix A from equality constraints
diff --git a/bofire/utils/torch_tools.py b/bofire/utils/torch_tools.py
index 9c7ef3772..beabcc98c 100644
--- a/bofire/utils/torch_tools.py
+++ b/bofire/utils/torch_tools.py
@@ -25,6 +25,7 @@
     CategoricalInput,
     CategoricalMolecularInput,
     ContinuousInput,
+    DiscreteInput,
     Input,
 )
 from bofire.data_models.molfeatures.api import AnyMolFeatures
@@ -56,7 +57,7 @@
 
 def get_linear_constraints(
     domain: Domain,
-    constraint: Union[Type[LinearEqualityConstraint], Type[LinearInequalityConstraint]],
+    constraint: Type[LinearEqualityConstraint] | Type[LinearInequalityConstraint],
     unit_scaled: bool = False,
 ) -> List[Tuple[Tensor, Tensor, float]]:
     """Converts linear constraints to the form required by BoTorch. For inequality constraints, this is A * x >= b (!).
@@ -71,6 +72,7 @@ def get_linear_constraints(
 
     """
     constraints = []
+    inputs = domain.inputs.get([ContinuousInput, DiscreteInput])
     for c in domain.constraints.get(constraint):
         indices = []
         coefficients = []
@@ -78,8 +80,8 @@ def get_linear_constraints(
         upper = []
         rhs = 0.0
         for i, featkey in enumerate(c.features):
-            idx = domain.inputs.get_keys(Input).index(featkey)
-            feat = domain.inputs.get_by_key(featkey)
+            idx = inputs.get_keys().index(featkey)
+            feat = inputs.get_by_key(featkey)
             if feat.is_fixed():
                 fixed = feat.fixed_value()
                 assert fixed is not None