Skip to content

metric

Compute a specified metric for the given tensor.

Parameters:

Name Type Description Default
x Tensor

The input tensor.

 required
metric_name str

The name of the metric to compute. Supported metrics: - "norm", "batch_norm", "l2_norm", "batch_l2_norm", "l1_norm", "batch_l1_norm" - "max", "batch_max", "min", "batch_min", "sum", "batch_sum", "prod", "batch_prod"

 required
*args

Additional arguments for specific metrics.

 ()
**kwargs

Additional arguments for specific metrics.

 ()

Returns:

Type Description
Tensor

The computed metric.

Raises:

Type Description
AssertionError

If the input tensor or metric name is None.
ValueError

If an unknown metric name is provided.
Source code in tinybig/koala/linear_algebra/metric.py 
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
def metric(
    x: torch.Tensor,
    metric_name: str,
    *args, **kwargs
):
    """
        Compute a specified metric for the given tensor.

        Parameters
        ----------
        x : torch.Tensor
            The input tensor.
        metric_name : str
            The name of the metric to compute. Supported metrics:
            - "norm", "batch_norm", "l2_norm", "batch_l2_norm", "l1_norm", "batch_l1_norm"
            - "max", "batch_max", "min", "batch_min", "sum", "batch_sum", "prod", "batch_prod"
        *args, **kwargs
            Additional arguments for specific metrics.

        Returns
        -------
        torch.Tensor
            The computed metric.

        Raises
        ------
        AssertionError
            If the input tensor or metric name is None.
        ValueError
            If an unknown metric name is provided.
        """

    assert x is not None and metric_name is not None

    match metric_name:
        case 'norm': return norm(x=x, *args, **kwargs)
        case 'batch_norm': return batch_norm(x=x, *args, **kwargs)
        case 'l2_norm': return l2_norm(x=x)
        case 'batch_l2_norm': return batch_l2_norm(x=x, *args, **kwargs)
        case 'l1_norm': return l1_norm(x=x)
        case 'batch_l1_norm': return batch_l1_norm(x=x, *args, **kwargs)
        case 'max': return max(x=x)
        case 'batch_max': return batch_max(x=x, *args, **kwargs)
        case 'min': return min(x=x)
        case 'batch_min': return batch_min(x=x, *args, **kwargs)
        case 'sum': return sum(x=x)
        case 'batch_sum': return batch_sum(x=x, *args, **kwargs)
        case 'prod': return prod(x=x)
        case 'batch_prod': return batch_prod(x=x, *args, **kwargs)
        case _: raise ValueError(f'Unknown metric name: {metric_name}...')

Compute the norm of a 1D tensor.

Parameters:

Name Type Description Default
x Tensor

The input 1D tensor.

 required
p Union[int, float, str, Any]

The norm order (e.g., 1, 2, "inf").

 required

Returns:

Type Description
Tensor

The computed norm.

Raises:

Type Description
AssertionError

If the input tensor is not 1D.
ValueError

If the nuclear norm is requested for a 1D tensor.
Source code in tinybig/koala/linear_algebra/metric.py 
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
def norm(x: torch.Tensor, p: Union[int, float, str, Any]):
    """
        Compute the norm of a 1D tensor.

        Parameters
        ----------
        x : torch.Tensor
            The input 1D tensor.
        p : Union[int, float, str, Any]
            The norm order (e.g., 1, 2, "inf").

        Returns
        -------
        torch.Tensor
            The computed norm.

        Raises
        ------
        AssertionError
            If the input tensor is not 1D.
        ValueError
            If the nuclear norm is requested for a 1D tensor.
        """
    assert x.ndim == 1
    if p == 'nuc':
        raise ValueError(f'the {p}-norm cannot be applied to 1d tensor inputs...')
    else:
        return torch.norm(x, p=p)

Compute the norm of a batch of tensors along a specified dimension.

Parameters:

Name Type Description Default
x Tensor

The input 2D tensor.

 required
p Union[int, float, str, Any]

The norm order (e.g., 1, 2, "nuc").

 required
dim int

The dimension along which to compute the norm. Default is 1.

 1

Returns:

Type Description
Tensor

The computed norms.

Raises:

Type Description
AssertionError

If the input tensor is not 2D or the dimension is invalid.
Source code in tinybig/koala/linear_algebra/metric.py 
 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
def batch_norm(x: torch.Tensor, p: Union[int, float, str, Any], dim: int = 1):
    """
        Compute the norm of a batch of tensors along a specified dimension.

        Parameters
        ----------
        x : torch.Tensor
            The input 2D tensor.
        p : Union[int, float, str, Any]
            The norm order (e.g., 1, 2, "nuc").
        dim : int, optional
            The dimension along which to compute the norm. Default is 1.

        Returns
        -------
        torch.Tensor
            The computed norms.

        Raises
        ------
        AssertionError
            If the input tensor is not 2D or the dimension is invalid.
        """
    assert x.ndim == 2 and dim in [0, 1, None]
    if p == 'nuc':
        return torch.norm(x, p='nuc')
    else:
        return torch.norm(x, p=p, dim=dim)

Compute the L1 norm of a 1D tensor.

Parameters:

Name Type Description Default
x Tensor

The input 1D tensor.

 required

Returns:

Type Description
Tensor

The L1 norm.
Source code in tinybig/koala/linear_algebra/metric.py 
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
def l1_norm(x: torch.Tensor):
    """
        Compute the L1 norm of a 1D tensor.

        Parameters
        ----------
        x : torch.Tensor
            The input 1D tensor.

        Returns
        -------
        torch.Tensor
            The L1 norm.
        """
    return norm(x, p=1)

Compute the L1 norm of a batch of tensors along a specified dimension.

Parameters:

Name Type Description Default
x Tensor

The input 2D tensor.

 required
dim int

The dimension along which to compute the L1 norm. Default is 1.

 1

Returns:

Type Description
Tensor

The L1 norms.
Source code in tinybig/koala/linear_algebra/metric.py 
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
def batch_l1_norm(x: torch.Tensor, dim: int = 1):
    """
        Compute the L1 norm of a batch of tensors along a specified dimension.

        Parameters
        ----------
        x : torch.Tensor
            The input 2D tensor.
        dim : int, optional
            The dimension along which to compute the L1 norm. Default is 1.

        Returns
        -------
        torch.Tensor
            The L1 norms.
        """
    return batch_norm(x, p=1, dim=dim)

Compute the L2 norm of a 1D tensor.

Parameters:

Name Type Description Default
x Tensor

The input 1D tensor.

 required

Returns:

Type Description
Tensor

The L2 norm.
Source code in tinybig/koala/linear_algebra/metric.py 
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
def l2_norm(x: torch.Tensor):
    """
        Compute the L2 norm of a 1D tensor.

        Parameters
        ----------
        x : torch.Tensor
            The input 1D tensor.

        Returns
        -------
        torch.Tensor
            The L2 norm.
        """
    return norm(x, p=2)

Compute the L2 norm of a batch of tensors along a specified dimension.

Parameters:

Name Type Description Default
x Tensor

The input 2D tensor.

 required
dim int

The dimension along which to compute the L2 norm. Default is 1.

 1

Returns:

Type Description
Tensor

The L2 norms.
Source code in tinybig/koala/linear_algebra/metric.py 
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
def batch_l2_norm(x: torch.Tensor, dim: int = 1):
    """
        Compute the L2 norm of a batch of tensors along a specified dimension.

        Parameters
        ----------
        x : torch.Tensor
            The input 2D tensor.
        dim : int, optional
            The dimension along which to compute the L2 norm. Default is 1.

        Returns
        -------
        torch.Tensor
            The L2 norms.
        """
    return batch_norm(x, p=2, dim=dim)

Compute the sum of elements in a 1D tensor.

Parameters:

Name Type Description Default
x Tensor

The input 1D tensor.

 required

Returns:

Type Description
Tensor

The sum of elements.
Source code in tinybig/koala/linear_algebra/metric.py 
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
def sum(x: torch.Tensor):
    """
        Compute the sum of elements in a 1D tensor.

        Parameters
        ----------
        x : torch.Tensor
            The input 1D tensor.

        Returns
        -------
        torch.Tensor
            The sum of elements.
        """
    assert x.ndim == 1
    return torch.sum(x)

Compute the sum of elements along a specified dimension in a batch of tensors.

Parameters:

Name Type Description Default
x Tensor

The input 2D tensor.

 required
dim int

The dimension along which to compute the sum. Default is 1.

 1

Returns:

Type Description
Tensor

The sums.
Source code in tinybig/koala/linear_algebra/metric.py 
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
def batch_sum(x: torch.Tensor, dim: int = 1):
    """
        Compute the sum of elements along a specified dimension in a batch of tensors.

        Parameters
        ----------
        x : torch.Tensor
            The input 2D tensor.
        dim : int, optional
            The dimension along which to compute the sum. Default is 1.

        Returns
        -------
        torch.Tensor
            The sums.
        """
    assert x.ndim == 2 and dim in [0, 1]
    return torch.sum(x, dim=dim)

Compute the product of elements in a 1D tensor.

Parameters:

Name Type Description Default
x Tensor

The input 1D tensor.

 required

Returns:

Type Description
Tensor

The product of elements.
Source code in tinybig/koala/linear_algebra/metric.py 
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
def prod(x: torch.Tensor):
    """
        Compute the product of elements in a 1D tensor.

        Parameters
        ----------
        x : torch.Tensor
            The input 1D tensor.

        Returns
        -------
        torch.Tensor
            The product of elements.
        """
    assert x.ndim == 1
    return torch.prod(x)

Compute the product of elements along a specified dimension in a batch of tensors.

Parameters:

Name Type Description Default
x Tensor

The input 2D tensor.

 required
dim int

The dimension along which to compute the product. Default is 1.

 1

Returns:

Type Description
Tensor

The products.
Source code in tinybig/koala/linear_algebra/metric.py 
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
def batch_prod(x: torch.Tensor, dim: int = 1):
    """
        Compute the product of elements along a specified dimension in a batch of tensors.

        Parameters
        ----------
        x : torch.Tensor
            The input 2D tensor.
        dim : int, optional
            The dimension along which to compute the product. Default is 1.

        Returns
        -------
        torch.Tensor
            The products.
        """
    assert x.ndim == 2 and dim in [0, 1]
    return torch.prod(x, dim=dim)

Compute the maximum value in a 1D tensor.

Parameters:

Name Type Description Default
x Tensor

The input 1D tensor.

 required

Returns:

Type Description
Tensor

The maximum value.
Source code in tinybig/koala/linear_algebra/metric.py 
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
def max(x: torch.Tensor):
    """
        Compute the maximum value in a 1D tensor.

        Parameters
        ----------
        x : torch.Tensor
            The input 1D tensor.

        Returns
        -------
        torch.Tensor
            The maximum value.
        """
    assert x.ndim == 1
    return torch.max(x)

Compute the maximum values along a specified dimension in a batch of tensors.

Parameters:

Name Type Description Default
x Tensor

The input 2D tensor.

 required
dim int

The dimension along which to compute the maximum values. Default is 1.

 1

Returns:

Type Description
Tensor

The maximum values.
Source code in tinybig/koala/linear_algebra/metric.py 
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
def batch_max(x: torch.Tensor, dim: int = 1):
    """
        Compute the maximum values along a specified dimension in a batch of tensors.

        Parameters
        ----------
        x : torch.Tensor
            The input 2D tensor.
        dim : int, optional
            The dimension along which to compute the maximum values. Default is 1.

        Returns
        -------
        torch.Tensor
            The maximum values.
        """
    assert x.ndim == 2 and dim in [0, 1]
    return torch.max(x, dim=dim).values

Compute the minimum value in a 1D tensor.

Parameters:

Name Type Description Default
x Tensor

The input 1D tensor.

 required

Returns:

Type Description
Tensor

The minimum value.
Source code in tinybig/koala/linear_algebra/metric.py 
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
def min(x: torch.Tensor):
    """
        Compute the minimum value in a 1D tensor.

        Parameters
        ----------
        x : torch.Tensor
            The input 1D tensor.

        Returns
        -------
        torch.Tensor
            The minimum value.
        """
    assert x.ndim == 1
    return torch.min(x)

Compute the minimum values along a specified dimension in a batch of tensors.

Parameters:

Name Type Description Default
x Tensor

The input 2D tensor.

 required
dim int

The dimension along which to compute the minimum values. Default is 1.

 1

Returns:

Type Description
Tensor

The minimum values.
Source code in tinybig/koala/linear_algebra/metric.py 
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
def batch_min(x: torch.Tensor, dim: int = 1):
    """
        Compute the minimum values along a specified dimension in a batch of tensors.

        Parameters
        ----------
        x : torch.Tensor
            The input 2D tensor.
        dim : int, optional
            The dimension along which to compute the minimum values. Default is 1.

        Returns
        -------
        torch.Tensor
            The minimum values.
        """
    assert x.ndim == 2 and dim in [0, 1]
    return torch.min(x, dim=dim).values