1import numpy as np
2
3import pyxu.abc as pxa
4import pyxu.info.ptype as pxt
5import pyxu.operator.linop.base as pxlb
6import pyxu.util as pxu
7
8__all__ = [
9 # Universal functions -----------------------------------------------------
10 "Sin",
11 "Cos",
12 "Tan",
13 "ArcSin",
14 "ArcCos",
15 "ArcTan",
16 "Sinh",
17 "Cosh",
18 "Tanh",
19 "ArcSinh",
20 "ArcCosh",
21 "ArcTanh",
22 "Exp",
23 "Log",
24 "Clip",
25 "Sqrt",
26 "Cbrt",
27 "Square",
28 "Abs",
29 "Sign",
30 "Gaussian",
31 "Sigmoid",
32 "SoftPlus",
33 "LeakyReLU",
34 "ReLU",
35 "SiLU",
36]
37
38
39# Trigonometric Functions =====================================================
[docs]
40class Sin(pxa.DiffMap):
41 r"""
42 Trigonometric sine, element-wise.
43
44 Notes
45 -----
46 * :math:`f(x) = \sin(x)`
47 * :math:`f'(x) = \cos(x)`
48 * :math:`\vert f(x) - f(y) \vert \le L \vert x - y \vert`, with Lipschitz constant :math:`L = 1`.
49
50 (Reason: :math:`\vert f'(x) \vert` is bounded by :math:`L` at :math:`x = k \pi, \, k \in \mathbb{Z}`.)
51 * :math:`\vert f'(x) - f'(y) \vert \le \partial L \vert x - y \vert`, with diff-Lipschitz constant :math:`\partial L
52 = 1`.
53
54 (Reason: :math:`\vert f''(x) \vert` is bounded by :math:`\partial L` at :math:`x = (2k + 1) \frac{\pi}{2}, \, k
55 \in \mathbb{Z}`.)
56 """
57
58 def __init__(self, dim_shape: pxt.NDArrayShape):
59 super().__init__(
60 dim_shape=dim_shape,
61 codim_shape=dim_shape,
62 )
63 self.lipschitz = 1
64 self.diff_lipschitz = 1
65
[docs]
66 def apply(self, arr: pxt.NDArray) -> pxt.NDArray:
67 xp = pxu.get_array_module(arr)
68 return xp.sin(arr)
69
[docs]
70 def jacobian(self, arr: pxt.NDArray) -> pxt.OpT:
71 xp = pxu.get_array_module(arr)
72 return pxlb.DiagonalOp(xp.cos(arr))
73
74
[docs]
75class Cos(pxa.DiffMap):
76 r"""
77 Trigonometric cosine, element-wise.
78
79 Notes
80 -----
81 * :math:`f(x) = \cos(x)`
82 * :math:`f'(x) = -\sin(x)`
83 * :math:`\vert f(x) - f(y) \vert \le L \vert x - y \vert`, with Lipschitz constant :math:`L = 1`.
84
85 (Reason: :math:`\vert f'(x) \vert` is bounded by :math:`L` at :math:`x = (2k + 1) \frac{\pi}{2}, \, k \in
86 \mathbb{Z}`.)
87 * :math:`\vert f'(x) - f'(y) \vert \le \partial L \vert x - y \vert`, with diff-Lipschitz constant :math:`\partial L
88 = 1`.
89
90 (Reason: :math:`\vert f''(x) \vert` is bounded by :math:`\partial L` at :math:`x = k \pi, \, k \in \mathbb{Z}`.)
91 """
92
93 def __init__(self, dim_shape: pxt.NDArrayShape):
94 super().__init__(
95 dim_shape=dim_shape,
96 codim_shape=dim_shape,
97 )
98 self.lipschitz = 1
99 self.diff_lipschitz = 1
100
[docs]
101 def apply(self, arr: pxt.NDArray) -> pxt.NDArray:
102 xp = pxu.get_array_module(arr)
103 return xp.cos(arr)
104
[docs]
105 def jacobian(self, arr: pxt.NDArray) -> pxt.OpT:
106 xp = pxu.get_array_module(arr)
107 return pxlb.DiagonalOp(-xp.sin(arr))
108
109
[docs]
110class Tan(pxa.DiffMap):
111 r"""
112 Trigonometric tangent, element-wise.
113
114 Notes
115 -----
116 * :math:`f(x) = \tan(x)`
117 * :math:`f'(x) = \cos^{-2}(x)`
118 * :math:`\vert f(x) - f(y) \vert \le L \vert x - y \vert`, with Lipschitz constant :math:`L = \infty`.
119
120 (Reason: :math:`f'(x)` is unbounded on :math:`\text{dom}(f) = [-\pi, \pi]`.)
121 * :math:`\vert f'(x) - f'(y) \vert \le \partial L \vert x - y \vert`, with diff-Lipschitz constant :math:`\partial L
122 = \infty`.
123
124 (Reason: :math:`f''(x)` is unbounded on :math:`\text{dom}(f) = [-\pi, \pi]`.)
125 """
126
127 def __init__(self, dim_shape: pxt.NDArrayShape):
128 super().__init__(
129 dim_shape=dim_shape,
130 codim_shape=dim_shape,
131 )
132 self.lipschitz = np.inf
133 self.diff_lipschitz = np.inf
134
[docs]
135 def apply(self, arr: pxt.NDArray) -> pxt.NDArray:
136 xp = pxu.get_array_module(arr)
137 return xp.tan(arr)
138
[docs]
139 def jacobian(self, arr: pxt.NDArray) -> pxt.OpT:
140 xp = pxu.get_array_module(arr)
141 v = xp.cos(arr)
142 v **= 2
143 return pxlb.DiagonalOp(1 / v)
144
145
[docs]
146class ArcSin(pxa.DiffMap):
147 r"""
148 Inverse sine, element-wise.
149
150 Notes
151 -----
152 * :math:`f(x) = \arcsin(x)`
153 * :math:`f'(x) = (1 - x^{2})^{-\frac{1}{2}}`
154 * :math:`\vert f(x) - f(y) \vert \le L \vert x - y \vert`, with Lipschitz constant :math:`L = \infty`.
155
156 (Reason: :math:`f'(x)` is unbounded on :math:`\text{dom}(f) = [-1, 1]`.)
157 * :math:`\vert f'(x) - f'(y) \vert \le \partial L \vert x - y \vert`, with diff-Lipschitz constant :math:`\partial L
158 = \infty`.
159
160 (Reason: :math:`f''(x)` is unbounded on :math:`\text{dom}(f) = [-1, 1]`.)
161 """
162
163 def __init__(self, dim_shape: pxt.NDArrayShape):
164 super().__init__(
165 dim_shape=dim_shape,
166 codim_shape=dim_shape,
167 )
168 self.lipschitz = np.inf
169 self.diff_lipschitz = np.inf
170
[docs]
171 def apply(self, arr: pxt.NDArray) -> pxt.NDArray:
172 xp = pxu.get_array_module(arr)
173 return xp.arcsin(arr)
174
[docs]
175 def jacobian(self, arr: pxt.NDArray) -> pxt.OpT:
176 xp = pxu.get_array_module(arr)
177 v = arr**2
178 v *= -1
179 v += 1
180 xp.sqrt(v, out=v)
181 return pxlb.DiagonalOp(1 / v)
182
183
[docs]
184class ArcCos(pxa.DiffMap):
185 r"""
186 Inverse cosine, element-wise.
187
188 Notes
189 -----
190 * :math:`f(x) = \arccos(x)`
191 * :math:`f'(x) = -(1 - x^{2})^{-\frac{1}{2}}`
192 * :math:`\vert f(x) - f(y) \vert \le L \vert x - y \vert`, with Lipschitz constant :math:`L = \infty`.
193
194 (Reason: :math:`f'(x)` is unbounded on :math:`\text{dom}(f) = [-1, 1]`.)
195 * :math:`\vert f'(x) - f'(y) \vert \le \partial L \vert x - y \vert`, with diff-Lipschitz constant :math:`\partial L
196 = \infty`.
197
198 (Reason: :math:`f''(x)` is unbounded on :math:`\text{dom}(f) = [-1, 1]`.)
199 """
200
201 def __init__(self, dim_shape: pxt.NDArrayShape):
202 super().__init__(
203 dim_shape=dim_shape,
204 codim_shape=dim_shape,
205 )
206 self.lipschitz = np.inf
207 self.diff_lipschitz = np.inf
208
[docs]
209 def apply(self, arr: pxt.NDArray) -> pxt.NDArray:
210 xp = pxu.get_array_module(arr)
211 return xp.arccos(arr)
212
[docs]
213 def jacobian(self, arr: pxt.NDArray) -> pxt.OpT:
214 xp = pxu.get_array_module(arr)
215 v = arr**2
216 v *= -1
217 v += 1
218 xp.sqrt(v, out=v)
219 return pxlb.DiagonalOp(-1 / v)
220
221
[docs]
222class ArcTan(pxa.DiffMap):
223 r"""
224 Inverse tangent, element-wise.
225
226 Notes
227 -----
228 * :math:`f(x) = \arctan(x)`
229 * :math:`f'(x) = (1 + x^{2})^{-1}`
230 * :math:`\vert f(x) - f(y) \vert \le L \vert x - y \vert`, with Lipschitz constant :math:`L = 1`.
231
232 (Reason: :math:`\vert f'(x) \vert` is bounded by :math:`L` at :math:`x = 0`.)
233 * :math:`\vert f'(x) - f'(y) \vert \le \partial L \vert x - y \vert`, with diff-Lipschitz constant :math:`\partial L
234 = 3 \sqrt{3} / 8`.
235
236 (Reason: :math:`\vert f''(x) \vert` is bounded by :math:`\partial L` at :math:`x = \pm \frac{1}{\sqrt{3}}`.)
237 """
238
239 def __init__(self, dim_shape: pxt.NDArrayShape):
240 super().__init__(
241 dim_shape=dim_shape,
242 codim_shape=dim_shape,
243 )
244 self.lipschitz = 1
245 self.diff_lipschitz = 3 * np.sqrt(3) / 8
246 # max_{x \in R} |arctan''(x)|
247 # = max_{x \in R} |2x / (1+x^2)^2|
248 # = 3 \sqrt(3) / 8 [at x = +- 1/\sqrt(3)]
249
[docs]
250 def apply(self, arr: pxt.NDArray) -> pxt.NDArray:
251 xp = pxu.get_array_module(arr)
252 return xp.arctan(arr)
253
[docs]
254 def jacobian(self, arr: pxt.NDArray) -> pxt.OpT:
255 v = arr**2
256 v += 1
257 return pxlb.DiagonalOp(1 / v)
258
259
260# Hyperbolic Functions ========================================================
[docs]
261class Sinh(pxa.DiffMap):
262 r"""
263 Hyperbolic sine, element-wise.
264
265 Notes
266 -----
267 * :math:`f(x) = \sinh(x)`
268 * :math:`f'(x) = \cosh(x)`
269 * :math:`\vert f(x) - f(y) \vert \le L \vert x - y \vert`, with Lipschitz constant :math:`L = \infty`.
270
271 (Reason: :math:`f'(x)` is unbounded on :math:`\text{dom}(f) = \mathbb{R}`.)
272 * :math:`\vert f'(x) - f'(y) \vert \le \partial L \vert x - y \vert`, with diff-Lipschitz constant :math:`\partial L
273 = \infty`.
274
275 (Reason: :math:`f''(x)` is unbounded on :math:`\text{dom}(f) = \mathbb{R}`.)
276 """
277
278 def __init__(self, dim_shape: pxt.NDArrayShape):
279 super().__init__(
280 dim_shape=dim_shape,
281 codim_shape=dim_shape,
282 )
283 self.lipschitz = np.inf
284 self.diff_lipschitz = np.inf
285
[docs]
286 def apply(self, arr: pxt.NDArray) -> pxt.NDArray:
287 xp = pxu.get_array_module(arr)
288 return xp.sinh(arr)
289
[docs]
290 def jacobian(self, arr: pxt.NDArray) -> pxt.OpT:
291 xp = pxu.get_array_module(arr)
292 return pxlb.DiagonalOp(xp.cosh(arr))
293
294
[docs]
295class Cosh(pxa.DiffMap):
296 r"""
297 Hyperbolic cosine, element-wise.
298
299 Notes
300 -----
301 * :math:`f(x) = \cosh(x)`
302 * :math:`f'(x) = \sinh(x)`
303 * :math:`\vert f(x) - f(y) \vert \le L \vert x - y \vert`, with Lipschitz constant :math:`L = \infty`.
304
305 (Reason: :math:`f'(x)` is unbounded on :math:`\text{dom}(f) = \mathbb{R}`.)
306 * :math:`\vert f'(x) - f'(y) \vert \le \partial L \vert x - y \vert`, with diff-Lipschitz constant :math:`\partial L
307 = \infty`.
308
309 (Reason: :math:`f''(x)` is unbounded on :math:`\text{dom}(f) = \mathbb{R}`.)
310 """
311
312 def __init__(self, dim_shape: pxt.NDArrayShape):
313 super().__init__(
314 dim_shape=dim_shape,
315 codim_shape=dim_shape,
316 )
317 self.lipschitz = np.inf
318 self.diff_lipschitz = np.inf
319
[docs]
320 def apply(self, arr: pxt.NDArray) -> pxt.NDArray:
321 xp = pxu.get_array_module(arr)
322 return xp.cosh(arr)
323
[docs]
324 def jacobian(self, arr: pxt.NDArray) -> pxt.OpT:
325 xp = pxu.get_array_module(arr)
326 return pxlb.DiagonalOp(xp.sinh(arr))
327
328
[docs]
329class Tanh(pxa.DiffMap):
330 r"""
331 Hyperbolic tangent, element-wise.
332
333 Notes
334 -----
335 * :math:`f(x) = \tanh(x)`
336 * :math:`f'(x) = 1 - \tanh^{2}(x)`
337 * :math:`\vert f(x) - f(y) \vert \le L \vert x - y \vert`, with Lipschitz constant :math:`L = 1`.
338
339 (Reason: :math:`\vert f'(x) \vert` is bounded by :math:`L` at :math:`x = 0`.)
340 * :math:`\vert f'(x) - f'(y) \vert \le \partial L \vert x - y \vert`, with diff-Lipschitz constant :math:`\partial L
341 = 4 / 3 \sqrt{3}`.
342
343 (Reason: :math:`\vert f''(x) \vert` is bounded by :math:`\partial L` at :math:`x = \frac{1}{2} \ln(2 \pm
344 \sqrt{3})`.
345 """
346
347 def __init__(self, dim_shape: pxt.NDArrayShape):
348 super().__init__(
349 dim_shape=dim_shape,
350 codim_shape=dim_shape,
351 )
352 self.lipschitz = 1
353 self.diff_lipschitz = 4 / (3 * np.sqrt(3))
354 # max_{x \in R} |tanh''(x)|
355 # = max_{x \in R} |-2 tanh(x) [1 - tanh(x)^2|
356 # = 4 / (3 \sqrt(3)) [at x = ln(2 +- \sqrt(3))]
357
[docs]
358 def apply(self, arr: pxt.NDArray) -> pxt.NDArray:
359 xp = pxu.get_array_module(arr)
360 return xp.tanh(arr)
361
[docs]
362 def jacobian(self, arr: pxt.NDArray) -> pxt.OpT:
363 v = self.apply(arr)
364 v = v**2
365 v *= -1
366 v += 1
367 return pxlb.DiagonalOp(v)
368
369
[docs]
370class ArcSinh(pxa.DiffMap):
371 r"""
372 Inverse hyperbolic sine, element-wise.
373
374 Notes
375 -----
376 * :math:`f(x) = \sinh^{-1}(x) = \ln(x + \sqrt{x^{2} + 1})`
377 * :math:`f'(x) = (x^{2} + 1)^{-\frac{1}{2}}`
378 * :math:`\vert f(x) - f(y) \vert \le L \vert x - y \vert`, with Lipschitz constant :math:`L = 1`.
379
380 (Reason: :math:`\vert f'(x) \vert` is bounded by :math:`L` at :math:`x = 0`.)
381 * :math:`\vert f'(x) - f'(y) \vert \le \partial L \vert x - y \vert`, with diff-Lipschitz constant :math:`\partial L
382 = \frac{2}{3 \sqrt{3}}`.
383
384 (Reason: :math:`\vert f''(x) \vert` is bounded by :math:`\partial L` at :math:`x = \pm \frac{1}{\sqrt{2}}`.)
385 """
386
387 def __init__(self, dim_shape: pxt.NDArrayShape):
388 super().__init__(
389 dim_shape=dim_shape,
390 codim_shape=dim_shape,
391 )
392 self.lipschitz = 1
393 self.diff_lipschitz = 2 / (3 * np.sqrt(3))
394 # max_{x \in R} |arcsinh''(x)|
395 # = max_{x \in R} |-x (x^2 + 1)^{-3/2}|
396 # = 2 / (3 \sqrt(3)) [at x = += 1 / \sqrt(2)]
397
[docs]
398 def apply(self, arr: pxt.NDArray) -> pxt.NDArray:
399 xp = pxu.get_array_module(arr)
400 return xp.arcsinh(arr)
401
[docs]
402 def jacobian(self, arr: pxt.NDArray) -> pxt.OpT:
403 xp = pxu.get_array_module(arr)
404 v = arr**2
405 v += 1
406 xp.sqrt(v, out=v)
407 return pxlb.DiagonalOp(1 / v)
408
409
[docs]
410class ArcCosh(pxa.DiffMap):
411 r"""
412 Inverse hyperbolic cosine, element-wise.
413
414 Notes
415 -----
416 * :math:`f(x) = \cosh^{-1}(x) = \ln(x + \sqrt{x^{2} - 1})`
417 * :math:`f'(x) = (x^{2} - 1)^{-\frac{1}{2}}`
418 * :math:`\vert f(x) - f(y) \vert \le L \vert x - y \vert`, with Lipschitz constant :math:`L = \infty`.
419
420 (Reason: :math:`f'(x)` is unbounded on :math:`\text{dom}(f) = [1, \infty[`.)
421 * :math:`\vert f'(x) - f'(y) \vert \le \partial L \vert x - y \vert`, with diff-Lipschitz constant :math:`\partial L
422 = \infty`.
423
424 (Reason: :math:`f''(x)` is unbounded on :math:`\text{dom}(f) = [1, \infty[`.)
425 """
426
427 def __init__(self, dim_shape: pxt.NDArrayShape):
428 super().__init__(
429 dim_shape=dim_shape,
430 codim_shape=dim_shape,
431 )
432 self.lipschitz = np.inf
433 self.diff_lipschitz = np.inf
434
[docs]
435 def apply(self, arr: pxt.NDArray) -> pxt.NDArray:
436 xp = pxu.get_array_module(arr)
437 return xp.arccosh(arr)
438
[docs]
439 def jacobian(self, arr: pxt.NDArray) -> pxt.OpT:
440 xp = pxu.get_array_module(arr)
441 v = arr**2
442 v -= 1
443 xp.sqrt(v, out=v)
444 return pxlb.DiagonalOp(1 / v)
445
446
[docs]
447class ArcTanh(pxa.DiffMap):
448 r"""
449 Inverse hyperbolic tangent, element-wise.
450
451 Notes
452 -----
453 * :math:`f(x) = \tanh^{-1}(x) = \frac{1}{2}\ln\left(\frac{1+x}{1-x}\right)`
454 * :math:`f'(x) = (1 - x^{2})^{-1}`
455 * :math:`\vert f(x) - f(y) \vert \le L \vert x - y \vert`, with Lipschitz constant :math:`L = \infty`.
456
457 (Reason: :math:`f'(x)` is unbounded on :math:`\text{dom}(f) = [-1, 1]`.)
458 * :math:`\vert f'(x) - f'(y) \vert \le \partial L \vert x - y \vert`, with diff-Lipschitz constant :math:`\partial L
459 = \infty`.
460
461 (Reason: :math:`f''(x)` is unbounded on :math:`\text{dom}(f) = [-1, 1]`.)
462 """
463
464 def __init__(self, dim_shape: pxt.NDArrayShape):
465 super().__init__(
466 dim_shape=dim_shape,
467 codim_shape=dim_shape,
468 )
469 self.lipschitz = np.inf
470 self.diff_lipschitz = np.inf
471
[docs]
472 def apply(self, arr: pxt.NDArray) -> pxt.NDArray:
473 xp = pxu.get_array_module(arr)
474 return xp.arctanh(arr)
475
[docs]
476 def jacobian(self, arr: pxt.NDArray) -> pxt.OpT:
477 v = arr**2
478 v *= -1
479 v += 1
480 return pxlb.DiagonalOp(1 / v)
481
482
483# Exponential Functions =======================================================
[docs]
484class Exp(pxa.DiffMap):
485 r"""
486 Exponential, element-wise. (Default: base-E exponential.)
487
488 Notes
489 -----
490 * :math:`f_{b}(x) = b^{x}`
491 * :math:`f_{b}'(x) = b^{x} \ln(b)`
492 * :math:`\vert f_{b}(x) - f_{b}(y) \vert \le L \vert x - y \vert`, with Lipschitz constant :math:`L = \infty`.
493
494 (Reason: :math:`f_{b}'(x)` is unbounded on :math:`\text{dom}(f_{b}) = \mathbb{R}`.)
495 * :math:`\vert f_{b}'(x) - f_{b}'(y) \vert \le \partial L \vert x - y \vert`, with diff-Lipschitz constant
496 :math:`\partial L = \infty`.
497
498 (Reason: :math:`f_{b}''(x)` is unbounded on :math:`\text{dom}(f_{b}) = \mathbb{R}`.)
499 """
500
501 def __init__(self, dim_shape: pxt.NDArrayShape, base: pxt.Real = None):
502 super().__init__(
503 dim_shape=dim_shape,
504 codim_shape=dim_shape,
505 )
506 self.lipschitz = np.inf
507 self.diff_lipschitz = np.inf
508 self._base = base
509
[docs]
510 def apply(self, arr: pxt.NDArray) -> pxt.NDArray:
511 out = arr.copy()
512 if self._base is not None:
513 out *= np.log(float(self._base))
514
515 xp = pxu.get_array_module(arr)
516 xp.exp(out, out=out)
517 return out
518
[docs]
519 def jacobian(self, arr: pxt.NDArray) -> pxt.OpT:
520 v = self.apply(arr)
521 if self._base is not None:
522 v *= np.log(float(self._base))
523 return pxlb.DiagonalOp(v)
524
525
[docs]
526class Log(pxa.DiffMap):
527 r"""
528 Logarithm, element-wise. (Default: base-E logarithm.)
529
530 Notes
531 -----
532 * :math:`f_{b}(x) = \log_{b}(x)`
533 * :math:`f_{b}'(x) = x^{-1} / \ln(b)`
534 * :math:`\vert f_{b}(x) - f_{b}(y) \vert \le L \vert x - y \vert`, with Lipschitz constant :math:`L = \infty`.
535
536 (Reason: :math:`f_{b}'(x)` is unbounded on :math:`\text{dom}(f_{b}) = \mathbb{R}_{+}`.)
537 * :math:`\vert f_{b}'(x) - f_{b}'(y) \vert \le \partial L \vert x - y \vert`, with diff-Lipschitz constant
538 :math:`\partial L = \infty`.
539
540 (Reason: :math:`f_{b}''(x)` is unbounded on :math:`\text{dom}(f_{b}) = \mathbb{R}_{+}`.)
541 """
542
543 def __init__(self, dim_shape: pxt.NDArrayShape, base: pxt.Real = None):
544 super().__init__(
545 dim_shape=dim_shape,
546 codim_shape=dim_shape,
547 )
548 self.lipschitz = np.inf
549 self.diff_lipschitz = np.inf
550 self._base = base
551
[docs]
552 def apply(self, arr: pxt.NDArray) -> pxt.NDArray:
553 xp = pxu.get_array_module(arr)
554 out = xp.log(arr)
555 if self._base is not None:
556 out /= np.log(float(self._base))
557 return out
558
[docs]
559 def jacobian(self, arr: pxt.NDArray) -> pxt.OpT:
560 v = 1 / arr
561 if self._base is not None:
562 v /= np.log(float(self._base))
563 return pxlb.DiagonalOp(v)
564
565
566# Miscellaneous ===============================================================
[docs]
567class Clip(pxa.Map):
568 r"""
569 Clip (limit) values in an array, element-wise.
570
571 Notes
572 -----
573 * .. math::
574
575 f_{[a,b]}(x) =
576 \begin{cases}
577 a, & \text{if} \ x \leq a, \\
578 x, & a < x < b, \\
579 b, & \text{if} \ x \geq b.
580 \end{cases}
581 * :math:`\vert f(x) - f(y) \vert \le L \vert x - y \vert`, with Lipschitz constant :math:`L = 1`.
582 """
583
584 def __init__(
585 self,
586 dim_shape: pxt.NDArrayShape,
587 a_min: pxt.Real = None,
588 a_max: pxt.Real = None,
589 ):
590 super().__init__(
591 dim_shape=dim_shape,
592 codim_shape=dim_shape,
593 )
594 if (a_min is None) and (a_max is None):
595 raise ValueError("One of Parameter[a_min, a_max] must be specified.")
596 self._llim = a_min
597 self._ulim = a_max
598 self.lipschitz = 1
599
[docs]
600 def apply(self, arr: pxt.NDArray) -> pxt.NDArray:
601 xp = pxu.get_array_module(arr)
602 out = arr.copy()
603 xp.clip(
604 arr,
605 self._llim,
606 self._ulim,
607 out=out,
608 )
609 return out
610
611
[docs]
612class Sqrt(pxa.DiffMap):
613 r"""
614 Non-negative square-root, element-wise.
615
616 Notes
617 -----
618 * :math:`f(x) = \sqrt{x}`
619 * :math:`f'(x) = 1 / 2 \sqrt{x}`
620 * :math:`\vert f(x) - f(y) \vert \le L \vert x - y \vert`, with Lipschitz constant :math:`L = \infty`.
621
622 (Reason: :math:`f'(x)` is unbounded on :math:`\text{dom}(f) = \mathbb{R}_{+}`.)
623 * :math:`\vert f'(x) - f'(y) \vert \le \partial L \vert x - y \vert`, with diff-Lipschitz constant :math:`\partial L
624 = \infty`.
625
626 (Reason: :math:`f''(x)` is unbounded on :math:`\text{dom}(f) = \mathbb{R}_{+}`.)
627 """
628
629 def __init__(self, dim_shape: pxt.NDArrayShape):
630 super().__init__(
631 dim_shape=dim_shape,
632 codim_shape=dim_shape,
633 )
634 self.lipschitz = np.inf
635 self.diff_lipschitz = np.inf
636
[docs]
637 def apply(self, arr: pxt.NDArray) -> pxt.NDArray:
638 xp = pxu.get_array_module(arr)
639 return xp.sqrt(arr)
640
[docs]
641 def jacobian(self, arr: pxt.NDArray) -> pxt.OpT:
642 v = self.apply(arr)
643 v *= 2
644 return pxlb.DiagonalOp(1 / v)
645
646
[docs]
647class Cbrt(pxa.DiffMap):
648 r"""
649 Cube-root, element-wise.
650
651 Notes
652 -----
653 * :math:`f(x) = \sqrt[3]{x}`
654 * :math:`f'(x) = 1 / 3 \sqrt[3]{x^{2}}`
655 * :math:`\vert f(x) - f(y) \vert \le L \vert x - y \vert`, with Lipschitz constant :math:`L = \infty`.
656
657 (Reason: :math:`f'(x)` is unbounded on :math:`\text{dom}(f) = \mathbb{R}`.)
658 * :math:`\vert f'(x) - f'(y) \vert \le \partial L \vert x - y \vert`, with diff-Lipschitz constant :math:`\partial L
659 = \infty`.
660
661 (Reason: :math:`f''(x)` is unbounded on :math:`\text{dom}(f) = \mathbb{R}`.)
662 """
663
664 def __init__(self, dim_shape: pxt.NDArrayShape):
665 super().__init__(
666 dim_shape=dim_shape,
667 codim_shape=dim_shape,
668 )
669 self.lipschitz = np.inf
670 self.diff_lipschitz = np.inf
671
[docs]
672 def apply(self, arr: pxt.NDArray) -> pxt.NDArray:
673 xp = pxu.get_array_module(arr)
674 return xp.cbrt(arr)
675
[docs]
676 def jacobian(self, arr: pxt.NDArray) -> pxt.OpT:
677 v = self.apply(arr)
678 v **= 2
679 v *= 3
680 return pxlb.DiagonalOp(1 / v)
681
682
[docs]
683class Square(pxa.DiffMap):
684 r"""
685 Square, element-wise.
686
687 Notes
688 -----
689 * :math:`f(x) = x^{2}`
690 * :math:`f'(x) = 2 x`
691 * :math:`\vert f(x) - f(y) \vert \le L \vert x - y \vert`, with Lipschitz constant :math:`L = \infty`.
692
693 (Reason: :math:`f'(x)` is unbounded on :math:`\text{dom}(f) = \mathbb{R}`.)
694 * :math:`\vert f'(x) - f'(y) \vert \le \partial L \vert x - y \vert`, with diff-Lipschitz constant :math:`\partial L
695 = 2`.
696
697 (Reason: :math:`\vert f''(x) \vert` is bounded by :math:`\partial L` everywhere.)
698 """
699
700 def __init__(self, dim_shape: pxt.NDArrayShape):
701 super().__init__(
702 dim_shape=dim_shape,
703 codim_shape=dim_shape,
704 )
705 self.lipschitz = np.inf
706 self.diff_lipschitz = 2
707
[docs]
708 def apply(self, arr: pxt.NDArray) -> pxt.NDArray:
709 out = arr.copy()
710 out **= 2
711 return out
712
[docs]
713 def jacobian(self, arr: pxt.NDArray) -> pxt.OpT:
714 v = arr.copy()
715 v *= 2
716 return pxlb.DiagonalOp(v)
717
718
[docs]
719class Abs(pxa.Map):
720 r"""
721 Absolute value, element-wise.
722
723 Notes
724 -----
725 * :math:`f(x) = \vert x \vert`
726 * :math:`\vert f(x) - f(y) \vert \le L \vert x - y \vert`, with Lipschitz constant :math:`L = 1`.
727 """
728
729 def __init__(self, dim_shape: pxt.NDArrayShape):
730 super().__init__(
731 dim_shape=dim_shape,
732 codim_shape=dim_shape,
733 )
734 self.lipschitz = 1
735
[docs]
736 def apply(self, arr: pxt.NDArray) -> pxt.NDArray:
737 xp = pxu.get_array_module(arr)
738 return xp.abs(arr)
739
740
[docs]
741class Sign(pxa.Map):
742 r"""
743 Number sign indicator, element-wise.
744
745 Notes
746 -----
747 * :math:`f(x) = x / \vert x \vert`
748 * :math:`\vert f(x) - f(y) \vert \le L \vert x - y \vert`, with Lipschitz constant :math:`L = 2`.
749 """
750
751 def __init__(self, dim_shape: pxt.NDArrayShape):
752 super().__init__(
753 dim_shape=dim_shape,
754 codim_shape=dim_shape,
755 )
756
[docs]
757 def apply(self, arr: pxt.NDArray) -> pxt.NDArray:
758 xp = pxu.get_array_module(arr)
759 return xp.sign(arr)
760
761
762# Activation Functions ========================================================
[docs]
763class Gaussian(pxa.DiffMap):
764 r"""
765 Gaussian, element-wise.
766
767 Notes
768 -----
769 * :math:`f(x) = \exp(-x^{2})`
770 * :math:`f'(x) = -2 x \exp(-x^{2})`
771 * :math:`\vert f(x) - f(y) \vert \le L \vert x - y \vert`, with Lipschitz constant :math:`L = \sqrt{2 / e}`.
772
773 (Reason: :math:`\vert f'(x) \vert` is bounded by :math:`L` at :math:`x = \pm 1 / \sqrt{2}`.)
774 * :math:`\vert f'(x) - f'(y) \vert \le \partial L \vert x - y \vert`, with diff-Lipschitz constant :math:`\partial L
775 = 2`.
776
777 (Reason: :math:`\vert f''(x) \vert` is bounded by :math:`\partial L` at :math:`x = 0`.)
778 """
779
780 def __init__(self, dim_shape: pxt.NDArrayShape):
781 super().__init__(
782 dim_shape=dim_shape,
783 codim_shape=dim_shape,
784 )
785 self.lipschitz = np.sqrt(2 / np.e)
786 self.diff_lipschitz = 2
787
[docs]
788 def apply(self, arr: pxt.NDArray) -> pxt.NDArray:
789 xp = pxu.get_array_module(arr)
790 out = arr.copy()
791 out **= 2
792 out *= -1
793 xp.exp(out, out=out)
794 return out
795
[docs]
796 def jacobian(self, arr: pxt.NDArray) -> pxt.OpT:
797 v = self.apply(arr)
798 v *= -2
799 v *= arr
800 return pxlb.DiagonalOp(v)
801
802
[docs]
803class Sigmoid(pxa.DiffMap):
804 r"""
805 Sigmoid, element-wise.
806
807 Notes
808 -----
809 * :math:`f(x) = (1 + e^{-x})^{-1}`
810 * :math:`f'(x) = f(x) [ f(x) - 1 ]`
811 * :math:`\vert f(x) - f(y) \vert \le L \vert x - y \vert`, with Lipschitz constant :math:`L = 1 / 4`.
812
813 (Reason: :math:`\vert f'(x) \vert` is bounded by :math:`L` at :math:`x = 0`.)
814 * :math:`\vert f'(x) - f'(y) \vert \le \partial L \vert x - y \vert`, with diff-Lipschitz constant :math:`\partial L
815 = 1 / 6 \sqrt{3}`.
816
817 (Reason: :math:`\vert f''(x) \vert` is bounded by :math:`\partial L` at :math:`x = \ln(2 \pm \sqrt{3})`.)
818 """
819
820 def __init__(self, dim_shape: pxt.NDArrayShape):
821 super().__init__(
822 dim_shape=dim_shape,
823 codim_shape=dim_shape,
824 )
825 self.lipschitz = 1 / 4
826 self.diff_lipschitz = 1 / (6 * np.sqrt(3))
827
[docs]
828 def apply(self, arr: pxt.NDArray) -> pxt.NDArray:
829 xp = pxu.get_array_module(arr)
830 x = -arr
831 xp.exp(x, out=x)
832 x += 1
833 return 1 / x
834
[docs]
835 def jacobian(self, arr: pxt.NDArray) -> pxt.OpT:
836 x = self.apply(arr)
837 v = x.copy()
838 x -= 1
839 v *= x
840 return pxlb.DiagonalOp(v)
841
842
[docs]
843class SoftPlus(pxa.DiffMap):
844 r"""
845 Softplus operator.
846
847 Notes
848 -----
849 * :math:`f(x) = \ln(1 + e^{x})`
850 * :math:`f'(x) = (1 + e^{-x})^{-1}`
851 * :math:`\vert f(x) - f(y) \vert \le L \vert x - y \vert`, with Lipschitz constant :math:`L = 1`.
852
853 (Reason: :math:`\vert f'(x) \vert` is bounded by :math:`L` on :math:`\text{dom}(f) = \mathbb{R}`.)
854 * :math:`\vert f'(x) - f'(y) \vert \le \partial L \vert x - y \vert`, with diff-Lipschitz constant :math:`\partial L
855 = 1 / 4`.
856
857 (Reason: :math:`\vert f''(x) \vert` is bounded by :math:`\partial L` at :math:`x = 0`.)
858 """
859
860 def __init__(self, dim_shape: pxt.NDArrayShape):
861 super().__init__(
862 dim_shape=dim_shape,
863 codim_shape=dim_shape,
864 )
865 self.lipschitz = 1
866 self.diff_lipschitz = 1 / 4
867
[docs]
868 def apply(self, arr: pxt.NDArray) -> pxt.NDArray:
869 xp = pxu.get_array_module(arr)
870 return xp.logaddexp(0, arr)
871
[docs]
872 def jacobian(self, arr: pxt.NDArray) -> pxt.OpT:
873 f = Sigmoid(dim_shape=self.dim_shape)
874 v = f.apply(arr)
875 return pxlb.DiagonalOp(v)
876
877
[docs]
878class LeakyReLU(pxa.Map):
879 r"""
880 Leaky rectified linear unit, element-wise.
881
882 Notes
883 -----
884 * :math:`f(x) = x \left[\mathbb{1}_{\ge 0}(x) + \alpha \mathbb{1}_{< 0}(x)\right], \quad \alpha \ge 0`
885 * :math:`\vert f(x) - f(y) \vert \le L \vert x - y \vert`, with Lipschitz constant :math:`L = \max(1, \alpha)`.
886 """
887
888 def __init__(self, dim_shape: pxt.NDArrayShape, alpha: pxt.Real):
889 super().__init__(
890 dim_shape=dim_shape,
891 codim_shape=dim_shape,
892 )
893 self._alpha = float(alpha)
894 assert self._alpha >= 0
895 self.lipschitz = float(max(alpha, 1))
896
[docs]
897 def apply(self, arr: pxt.NDArray) -> pxt.NDArray:
898 xp = pxu.get_array_module(arr)
899 return xp.where(arr >= 0, arr, arr * self._alpha)
900
901
[docs]
902class ReLU(LeakyReLU):
903 r"""
904 Rectified linear unit, element-wise.
905
906 Notes
907 -----
908 * :math:`f(x) = \lfloor x \rfloor_{+}`
909 * :math:`\vert f(x) - f(y) \vert \le L \vert x - y \vert`, with Lipschitz constant :math:`L = 1`.
910 """
911
912 def __init__(self, dim_shape: pxt.NDArrayShape):
913 super().__init__(dim_shape=dim_shape, alpha=0)
914
915
[docs]
916class SiLU(pxa.DiffMap):
917 r"""
918 Sigmoid linear unit, element-wise.
919
920 Notes
921 -----
922 * :math:`f(x) = x / (1 + e^{-x})`
923 * :math:`f'(x) = (1 + e^{-x} + x e^{-x}) / (1 + e^{-x})^{2}`
924 * :math:`\vert f(x) - f(y) \vert \le L \vert x - y \vert`, with Lipschitz constant :math:`L = 1.1`.
925
926 (Reason: :math:`\vert f'(x) \vert` is bounded by :math:`L` at :math:`x \approx 2.4`.)
927 * :math:`\vert f'(x) - f'(y) \vert \le \partial L \vert x - y \vert`, with diff-Lipschitz constant :math:`\partial L
928 = 1 / 2`.
929
930 (Reason: :math:`\vert f''(x) \vert` is bounded by :math:`\partial L` at :math:`x = 0`.)
931 """
932
933 def __init__(self, dim_shape: pxt.NDArrayShape):
934 super().__init__(
935 dim_shape=dim_shape,
936 codim_shape=dim_shape,
937 )
938 self.lipschitz = 1.1
939 self.diff_lipschitz = 1 / 2
940
[docs]
941 def apply(self, arr: pxt.NDArray) -> pxt.NDArray:
942 f = Sigmoid(dim_shape=self.dim_shape)
943 out = f.apply(arr)
944 out *= arr
945 return out
946
[docs]
947 def jacobian(self, arr: pxt.NDArray) -> pxt.OpT:
948 f = Sigmoid(dim_shape=self.dim_shape)
949 xp = pxu.get_array_module(arr)
950 a = xp.exp(-arr)
951 a *= 1 + arr
952 a += 1
953 b = f.apply(arr)
954 b **= 2
955 return pxlb.DiagonalOp(a * b)
