Skip to content

BaseSampler

Methods and Attributes

Bases: ABC

Base class for samplers.

Parameters:

Name Type Description Default
energy_function BaseEnergyFunction

Energy function to sample from.

 required
dtype dtype

Data type to use for the computations.

 float32
device Union[str, device]

Device to run the computations on (e.g., "cpu" or "cuda").

 None
use_mixed_precision bool

Whether to use mixed precision for sampling operations.

 False

Methods:

Name Description
sample

Run the sampling process.
sample_chain

Run the sampling process.
_setup_diagnostics

Initialize the diagnostics dictionary.
to

Move sampler to specified device.
Source code in torchebm/core/base_sampler.py 
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
class BaseSampler(ABC):
    """
    Base class for samplers.

    Args:
        energy_function (BaseEnergyFunction): Energy function to sample from.
        dtype (torch.dtype): Data type to use for the computations.
        device (Union[str, torch.device]): Device to run the computations on (e.g., "cpu" or "cuda").
        use_mixed_precision (bool): Whether to use mixed precision for sampling operations.

    Methods:
        sample(x, dim, k_steps, n_samples, thin, return_trajectory, return_diagnostics): Run the sampling process.
        sample_chain(dim, k_steps, n_samples, thin, return_trajectory, return_diagnostics): Run the sampling process.
        _setup_diagnostics(): Initialize the diagnostics dictionary.
        to(device): Move sampler to specified device.
    """

    def __init__(
        self,
        energy_function: BaseEnergyFunction,
        dtype: torch.dtype = torch.float32,
        device: Optional[Union[str, torch.device]] = None,
        use_mixed_precision: bool = False,
    ):
        self.energy_function = energy_function
        self.dtype = dtype
        if isinstance(device, str):
            device = torch.device(device)
        self.device = device or torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.use_mixed_precision = use_mixed_precision

        # Check if mixed precision is available
        if self.use_mixed_precision:
            try:
                from torch.cuda.amp import autocast
                self.autocast_available = True
                # Ensure device is CUDA for mixed precision
                if not self.device.type.startswith('cuda'):
                    warnings.warn(
                        f"Mixed precision requested but device is {self.device}. "
                        f"Mixed precision requires CUDA. Falling back to full precision.",
                        UserWarning,
                    )
                    self.use_mixed_precision = False
                    self.autocast_available = False
            except ImportError:
                warnings.warn(
                    "Mixed precision requested but torch.cuda.amp not available. "
                    "Falling back to full precision. Requires PyTorch 1.6+.",
                    UserWarning,
                )
                self.use_mixed_precision = False
                self.autocast_available = False
        else:
            self.autocast_available = False

        self.schedulers: Dict[str, BaseScheduler] = {}

        # Ensure the energy function has matching precision settings
        if hasattr(self.energy_function, 'use_mixed_precision'):
            self.energy_function.use_mixed_precision = self.use_mixed_precision

    @abstractmethod
    def sample(
        self,
        x: Optional[torch.Tensor] = None,
        dim: int = 10,
        n_steps: int = 100,
        n_samples: int = 1,
        thin: int = 1,
        return_trajectory: bool = False,
        return_diagnostics: bool = False,
        *args,
        **kwargs,
    ) -> Union[torch.Tensor, Tuple[torch.Tensor, List[dict]]]:
        """
        Run the sampling process.

        Args:
            x: Initial state to start the sampling from.
            dim: Dimension of the state space.
            k_steps: Number of steps to take between samples.
            n_samples: Number of samples to generate.
            thin: Thinning factor (not supported yet).
            return_trajectory: Whether to return the trajectory of the samples.
            return_diagnostics: Whether to return the diagnostics of the sampling process.

        Returns:
            torch.Tensor: Samples from the sampler.
            List[dict]: Diagnostics of the sampling process.
        """
        raise NotImplementedError

    def register_scheduler(self, name: str, scheduler: BaseScheduler) -> None:
        """
        Register a parameter scheduler.

        Args:
            name: Name of the parameter to schedule
            scheduler: Scheduler instance to use
        """
        self.schedulers[name] = scheduler

    def get_schedulers(self) -> Dict[str, BaseScheduler]:
        """
        Get all registered schedulers.

        Returns:
            Dictionary mapping parameter names to their schedulers
        """
        return self.schedulers

    def get_scheduled_value(self, name: str) -> float:
        """
        Get current value for a scheduled parameter.

        Args:
            name: Name of the scheduled parameter

        Returns:
            Current value of the parameter

        Raises:
            KeyError: If no scheduler exists for the parameter
        """
        if name not in self.schedulers:
            raise KeyError(f"No scheduler registered for parameter '{name}'")
        return self.schedulers[name].get_value()

    def step_schedulers(self) -> Dict[str, float]:
        """
        Advance all schedulers by one step.

        Returns:
            Dictionary mapping parameter names to their updated values
        """
        return {name: scheduler.step() for name, scheduler in self.schedulers.items()}

    def reset_schedulers(self) -> None:
        """Reset all schedulers to their initial state."""
        for scheduler in self.schedulers.values():
            scheduler.reset()

    # @abstractmethod
    def _setup_diagnostics(self) -> dict:
        """
        Initialize the diagnostics dictionary.

            .. deprecated:: 1.0
               This method is deprecated and will be removed in a future version.
        """
        return {
            "energies": torch.empty(0, device=self.device, dtype=self.dtype),
            "acceptance_rate": torch.tensor(0.0, device=self.device, dtype=self.dtype),
        }
        # raise NotImplementedError

    def to(self, device: Union[str, torch.device], dtype: Optional[torch.dtype] = None) -> "BaseSampler":
        """
        Move sampler to the specified device and optionally change its dtype.

        Args:
            device: Target device for computations
            dtype: Optional data type to convert to

        Returns:
            The sampler instance moved to the specified device/dtype
        """
        if isinstance(device, str):
            device = torch.device(device)

        self.device = device

        if dtype is not None:
            self.dtype = dtype

        # Update mixed precision availability if device changed
        if self.use_mixed_precision and not self.device.type.startswith('cuda'):
            warnings.warn(
                f"Mixed precision active but moving to {self.device}. "
                f"Mixed precision requires CUDA. Disabling mixed precision.",
                UserWarning,
            )
            self.use_mixed_precision = False

        # Move energy function if it has a to method
        if hasattr(self.energy_function, "to") and callable(getattr(self.energy_function, "to")):
            self.energy_function = self.energy_function.to(device=self.device, dtype=self.dtype)

        return self

    def apply_mixed_precision(self, func):
        """
        Decorator to apply mixed precision context to a method.

        Args:
            func: Function to wrap with mixed precision

        Returns:
            Wrapped function with mixed precision support
        """
        def wrapper(*args, **kwargs):
            if self.use_mixed_precision and self.autocast_available:
                from torch.cuda.amp import autocast
                with autocast():
                    return func(*args, **kwargs)
            else:
                return func(*args, **kwargs)
        return wrapper

energy_function instance-attribute 

energy_function = energy_function

dtype instance-attribute 

dtype = dtype

device instance-attribute 

device = device or device('cuda' if is_available() else 'cpu')

use_mixed_precision instance-attribute 

use_mixed_precision = use_mixed_precision

autocast_available instance-attribute 

autocast_available = True

schedulers instance-attribute 

schedulers: Dict[str, BaseScheduler] = {}

sample abstractmethod 

sample(x: Optional[Tensor] = None, dim: int = 10, n_steps: int = 100, n_samples: int = 1, thin: int = 1, return_trajectory: bool = False, return_diagnostics: bool = False, *args, **kwargs) -> Union[torch.Tensor, Tuple[torch.Tensor, List[dict]]]

Run the sampling process.

Parameters:

Name Type Description Default
x Optional[Tensor]

Initial state to start the sampling from.

 None
dim int

Dimension of the state space.

 10
k_steps

Number of steps to take between samples.

 required
n_samples int

Number of samples to generate.

 1
thin int

Thinning factor (not supported yet).

 1
return_trajectory bool

Whether to return the trajectory of the samples.

 False
return_diagnostics bool

Whether to return the diagnostics of the sampling process.

 False

Returns:

Type Description
Union[Tensor, Tuple[Tensor, List[dict]]]

torch.Tensor: Samples from the sampler.
Union[Tensor, Tuple[Tensor, List[dict]]]

List[dict]: Diagnostics of the sampling process.
Source code in torchebm/core/base_sampler.py 
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
@abstractmethod
def sample(
    self,
    x: Optional[torch.Tensor] = None,
    dim: int = 10,
    n_steps: int = 100,
    n_samples: int = 1,
    thin: int = 1,
    return_trajectory: bool = False,
    return_diagnostics: bool = False,
    *args,
    **kwargs,
) -> Union[torch.Tensor, Tuple[torch.Tensor, List[dict]]]:
    """
    Run the sampling process.

    Args:
        x: Initial state to start the sampling from.
        dim: Dimension of the state space.
        k_steps: Number of steps to take between samples.
        n_samples: Number of samples to generate.
        thin: Thinning factor (not supported yet).
        return_trajectory: Whether to return the trajectory of the samples.
        return_diagnostics: Whether to return the diagnostics of the sampling process.

    Returns:
        torch.Tensor: Samples from the sampler.
        List[dict]: Diagnostics of the sampling process.
    """
    raise NotImplementedError

register_scheduler 

register_scheduler(name: str, scheduler: BaseScheduler) -> None

Register a parameter scheduler.

Parameters:

Name Type Description Default
name str

Name of the parameter to schedule

 required
scheduler BaseScheduler

Scheduler instance to use

 required
Source code in torchebm/core/base_sampler.py 
103
104
105
106
107
108
109
110
111
def register_scheduler(self, name: str, scheduler: BaseScheduler) -> None:
    """
    Register a parameter scheduler.

    Args:
        name: Name of the parameter to schedule
        scheduler: Scheduler instance to use
    """
    self.schedulers[name] = scheduler

get_schedulers 

get_schedulers() -> Dict[str, BaseScheduler]

Get all registered schedulers.

Returns:

Type Description
Dict[str, BaseScheduler]

Dictionary mapping parameter names to their schedulers
Source code in torchebm/core/base_sampler.py 
113
114
115
116
117
118
119
120
def get_schedulers(self) -> Dict[str, BaseScheduler]:
    """
    Get all registered schedulers.

    Returns:
        Dictionary mapping parameter names to their schedulers
    """
    return self.schedulers

get_scheduled_value 

get_scheduled_value(name: str) -> float

Get current value for a scheduled parameter.

Parameters:

Name Type Description Default
name str

Name of the scheduled parameter

 required

Returns:

Type Description
float

Current value of the parameter

Raises:

Type Description
KeyError

If no scheduler exists for the parameter
Source code in torchebm/core/base_sampler.py 
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
def get_scheduled_value(self, name: str) -> float:
    """
    Get current value for a scheduled parameter.

    Args:
        name: Name of the scheduled parameter

    Returns:
        Current value of the parameter

    Raises:
        KeyError: If no scheduler exists for the parameter
    """
    if name not in self.schedulers:
        raise KeyError(f"No scheduler registered for parameter '{name}'")
    return self.schedulers[name].get_value()

step_schedulers 

step_schedulers() -> Dict[str, float]

Advance all schedulers by one step.

Returns:

Type Description
Dict[str, float]

Dictionary mapping parameter names to their updated values
Source code in torchebm/core/base_sampler.py 
139
140
141
142
143
144
145
146
def step_schedulers(self) -> Dict[str, float]:
    """
    Advance all schedulers by one step.

    Returns:
        Dictionary mapping parameter names to their updated values
    """
    return {name: scheduler.step() for name, scheduler in self.schedulers.items()}

reset_schedulers 

reset_schedulers() -> None

Reset all schedulers to their initial state.

Source code in torchebm/core/base_sampler.py 
148
149
150
151
def reset_schedulers(self) -> None:
    """Reset all schedulers to their initial state."""
    for scheduler in self.schedulers.values():
        scheduler.reset()

_setup_diagnostics 

_setup_diagnostics() -> dict

Initialize the diagnostics dictionary.

1
2
.. deprecated:: 1.0
   This method is deprecated and will be removed in a future version.
Source code in torchebm/core/base_sampler.py 
154
155
156
157
158
159
160
161
162
163
164
def _setup_diagnostics(self) -> dict:
    """
    Initialize the diagnostics dictionary.

        .. deprecated:: 1.0
           This method is deprecated and will be removed in a future version.
    """
    return {
        "energies": torch.empty(0, device=self.device, dtype=self.dtype),
        "acceptance_rate": torch.tensor(0.0, device=self.device, dtype=self.dtype),
    }

to 

to(device: Union[str, device], dtype: Optional[dtype] = None) -> BaseSampler

Move sampler to the specified device and optionally change its dtype.

Parameters:

Name Type Description Default
device Union[str, device]

Target device for computations

 required
dtype Optional[dtype]

Optional data type to convert to

 None

Returns:

Type Description
BaseSampler

The sampler instance moved to the specified device/dtype
Source code in torchebm/core/base_sampler.py 
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
def to(self, device: Union[str, torch.device], dtype: Optional[torch.dtype] = None) -> "BaseSampler":
    """
    Move sampler to the specified device and optionally change its dtype.

    Args:
        device: Target device for computations
        dtype: Optional data type to convert to

    Returns:
        The sampler instance moved to the specified device/dtype
    """
    if isinstance(device, str):
        device = torch.device(device)

    self.device = device

    if dtype is not None:
        self.dtype = dtype

    # Update mixed precision availability if device changed
    if self.use_mixed_precision and not self.device.type.startswith('cuda'):
        warnings.warn(
            f"Mixed precision active but moving to {self.device}. "
            f"Mixed precision requires CUDA. Disabling mixed precision.",
            UserWarning,
        )
        self.use_mixed_precision = False

    # Move energy function if it has a to method
    if hasattr(self.energy_function, "to") and callable(getattr(self.energy_function, "to")):
        self.energy_function = self.energy_function.to(device=self.device, dtype=self.dtype)

    return self

apply_mixed_precision 

apply_mixed_precision(func)

Decorator to apply mixed precision context to a method.

Parameters:

Name Type Description Default
func

Function to wrap with mixed precision

 required

Returns:

Type Description

Wrapped function with mixed precision support
Source code in torchebm/core/base_sampler.py 
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
def apply_mixed_precision(self, func):
    """
    Decorator to apply mixed precision context to a method.

    Args:
        func: Function to wrap with mixed precision

    Returns:
        Wrapped function with mixed precision support
    """
    def wrapper(*args, **kwargs):
        if self.use_mixed_precision and self.autocast_available:
            from torch.cuda.amp import autocast
            with autocast():
                return func(*args, **kwargs)
        else:
            return func(*args, **kwargs)
    return wrapper