Skip to content

utils

utils ¤

LOG_CLAMP_MIN = -708.3964185322641 module-attribute ¤

csafelog = ComplexSafeLog.apply module-attribute ¤

safelog = SafeLog.apply module-attribute ¤

ComplexSafeLog ¤

Bases: Function

A numerically safe natural logarithm autograd function for complex inputs.

In the backward pass, it replaces the NaNs and infinities arising in the gradient with zeros.

Source code in cirkit/backend/torch/utils.py 
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
class ComplexSafeLog(autograd.Function):
    """A numerically safe natural logarithm autograd function for complex inputs.

    In the backward pass, it replaces the NaNs and infinities
    arising in the gradient with zeros.
    """

    @staticmethod
    def forward(x: Tensor) -> Tensor:  # pylint: disable=arguments-differ
        y = torch.log(x)
        y.real.clamp_min_(LOG_CLAMP_MIN)
        if torch.is_complex(y):
            y.imag.clamp_min_(LOG_CLAMP_MIN)
        return y

    @staticmethod
    def setup_context(  # pylint: disable=arguments-differ
        ctx: Any, inputs: tuple[Tensor, ...], output: Tensor
    ) -> None:
        (x,) = inputs
        ctx.save_for_backward(x)

    @staticmethod
    def backward(ctx: Any, grad_output: Tensor) -> Tensor:  # pylint: disable=arguments-differ
        (x,) = ctx.saved_tensors
        return torch.nan_to_num(grad_output / x.conj())

backward(ctx, grad_output) staticmethod ¤

Source code in cirkit/backend/torch/utils.py 
62
63
64
65
@staticmethod
def backward(ctx: Any, grad_output: Tensor) -> Tensor:  # pylint: disable=arguments-differ
    (x,) = ctx.saved_tensors
    return torch.nan_to_num(grad_output / x.conj())

forward(x) staticmethod ¤

Source code in cirkit/backend/torch/utils.py 
47
48
49
50
51
52
53
@staticmethod
def forward(x: Tensor) -> Tensor:  # pylint: disable=arguments-differ
    y = torch.log(x)
    y.real.clamp_min_(LOG_CLAMP_MIN)
    if torch.is_complex(y):
        y.imag.clamp_min_(LOG_CLAMP_MIN)
    return y

setup_context(ctx, inputs, output) staticmethod ¤

Source code in cirkit/backend/torch/utils.py 
55
56
57
58
59
60
@staticmethod
def setup_context(  # pylint: disable=arguments-differ
    ctx: Any, inputs: tuple[Tensor, ...], output: Tensor
) -> None:
    (x,) = inputs
    ctx.save_for_backward(x)

SafeLog ¤

Bases: Function

A numerically safe natural logarithm autograd function.

In the backward pass, it replaces the NaNs and infinities arising in the gradient with zeros.

Source code in cirkit/backend/torch/utils.py 
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
class SafeLog(autograd.Function):
    """A numerically safe natural logarithm autograd function.

    In the backward pass, it replaces the NaNs and infinities
    arising in the gradient with zeros.
    """

    @staticmethod
    def forward(x: Tensor) -> Tensor:  # pylint: disable=arguments-differ
        return torch.log(x).clamp_min(LOG_CLAMP_MIN)

    @staticmethod
    def setup_context(  # pylint: disable=arguments-differ
        ctx: Any, inputs: tuple[Tensor, ...], output: Tensor
    ) -> None:
        (x,) = inputs
        ctx.save_for_backward(x)

    @staticmethod
    def backward(ctx: Any, grad_output: Tensor) -> Tensor:  # pylint: disable=arguments-differ
        (x,) = ctx.saved_tensors
        return torch.nan_to_num(grad_output / x)

backward(ctx, grad_output) staticmethod ¤

Source code in cirkit/backend/torch/utils.py 
30
31
32
33
@staticmethod
def backward(ctx: Any, grad_output: Tensor) -> Tensor:  # pylint: disable=arguments-differ
    (x,) = ctx.saved_tensors
    return torch.nan_to_num(grad_output / x)

forward(x) staticmethod ¤

Source code in cirkit/backend/torch/utils.py 
19
20
21
@staticmethod
def forward(x: Tensor) -> Tensor:  # pylint: disable=arguments-differ
    return torch.log(x).clamp_min(LOG_CLAMP_MIN)

setup_context(ctx, inputs, output) staticmethod ¤

Source code in cirkit/backend/torch/utils.py 
23
24
25
26
27
28
@staticmethod
def setup_context(  # pylint: disable=arguments-differ
    ctx: Any, inputs: tuple[Tensor, ...], output: Tensor
) -> None:
    (x,) = inputs
    ctx.save_for_backward(x)

flatten_dims(x, /, *, dims) ¤

Flatten the given dims in the input.

If the dims are not continuous, they will be permuted and flattened to the position of the first element in dims.

Intended to be used as a helper for some torch functions that can only work on one dim.

Parameters:

Name Type Description Default
x Tensor

The tensor to be flattened.

 required
dims Sequence[int]

The dimensions to flatten along, expected to be sorted.

 required

Returns:

Name Type Description
Tensor Tensor

The flattened tensor.
Source code in cirkit/backend/torch/utils.py 
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
def flatten_dims(x: Tensor, /, *, dims: Sequence[int]) -> Tensor:
    """Flatten the given dims in the input.

    If the dims are not continuous, they will be permuted and flattened to the position of the \
    first element in dims.

    Intended to be used as a helper for some torch functions that can only work on one dim.

    Args:
        x: The tensor to be flattened.
        dims: The dimensions to flatten along, expected to be sorted.

    Returns:
        Tensor: The flattened tensor.
    """
    if not dims:  # When dims[0] does not work.
        return x

    start_dim, end_dim = dims[0], dims[0] + len(dims)
    # Note that for flatten, end_dim is inclusive.
    return x.movedim(tuple(dims), tuple(range(start_dim, end_dim))).flatten(start_dim, end_dim - 1)

unflatten_dims(x, /, *, dims, shape) ¤

Unflatten the first dim in dims in the input to get a given shape.

This is the inverse transformation of flatten_dims, provided a correspondimg shape.

Parameters:

Name Type Description Default
x Tensor

The tensor to be unflattened.

 required
dims Sequence[int]

The dimensions to unflatten to, should be the same as flatten_dims.

 required
shape Sequence[int]

The shape to unflatten to, can be either the shape for dims, or the whole shape for the output. If the latter, the shape will not be checked for consistency outside dims.

 required

Returns:

Name Type Description
Tensor Tensor

The unflattened tensor.
Source code in cirkit/backend/torch/utils.py 
 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
def unflatten_dims(x: Tensor, /, *, dims: Sequence[int], shape: Sequence[int]) -> Tensor:
    """Unflatten the first dim in dims in the input to get a given shape.

    This is the inverse transformation of flatten_dims, provided a correspondimg shape.

    Args:
        x (Tensor): The tensor to be unflattened.
        dims (Sequence[int]): The dimensions to unflatten to, should be the same as flatten_dims.
        shape (Sequence[int]): The shape to unflatten to, can be either the shape for dims, or the \
            whole shape for the output. If the latter, the shape will not be checked for \
            consistency outside dims.

    Returns:
        Tensor: The unflattened tensor.
    """
    if not dims:  # When dims[0] does not work.
        return x

    # We require dims to be sorted so that there's no ambiguation in how shape is interpreted,
    # unless the shape itself never causes ambiguation.
    assert all(s == 1 for s in shape) or all(
        l < r for l, r in itertools.pairwise(dims)
    ), "dims must be sorted for unflatten_dims."

    if len(shape) == x.ndim - 1 + len(dims):  # The shape is for whole output.
        shape = [shape[d] for d in dims]
    # The shape is now for dims.

    start_dim, end_dim = dims[0], dims[0] + len(dims)
    # TODO: x.unflatten is not typed, must use torch.unflatten for now.
    return torch.unflatten(x, dim=start_dim, sizes=shape).movedim(
        tuple(range(start_dim, end_dim)), tuple(dims)
    )