1
2
3
4
5
6
7
8
9
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
|
########
einx API
########
Main
----
.. autofunction:: einx.rearrange
.. autofunction:: einx.vmap_with_axis
.. autofunction:: einx.vmap
.. autofunction:: einx.reduce
.. autofunction:: einx.elementwise
.. autofunction:: einx.index
Reduction operations
--------------------
.. autofunction:: einx.sum
.. autofunction:: einx.mean
.. autofunction:: einx.var
.. autofunction:: einx.std
.. autofunction:: einx.prod
.. autofunction:: einx.count_nonzero
.. autofunction:: einx.any
.. autofunction:: einx.all
.. autofunction:: einx.max
.. autofunction:: einx.min
.. autofunction:: einx.logsumexp
Element-by-element operations
-----------------------------
.. autofunction:: einx.add
.. autofunction:: einx.subtract
.. autofunction:: einx.multiply
.. autofunction:: einx.true_divide
.. autofunction:: einx.floor_divide
.. autofunction:: einx.divide
.. autofunction:: einx.logical_and
.. autofunction:: einx.logical_or
.. autofunction:: einx.where
.. autofunction:: einx.less
.. autofunction:: einx.less_equal
.. autofunction:: einx.greater
.. autofunction:: einx.greater_equal
.. autofunction:: einx.equal
.. autofunction:: einx.not_equal
.. autofunction:: einx.maximum
.. autofunction:: einx.minimum
Indexing operations
-------------------
.. autofunction:: einx.get_at
.. autofunction:: einx.set_at
.. autofunction:: einx.add_at
.. autofunction:: einx.subtract_at
Miscellaneous operations
------------------------
.. autofunction:: einx.flip
.. autofunction:: einx.roll
.. autofunction:: einx.softmax
.. autofunction:: einx.log_softmax
.. autofunction:: einx.arange
General dot-product
-------------------
.. autofunction:: einx.dot
Deep Learning Modules
=====================
Haiku
-----
.. autoclass:: einx.nn.haiku.Linear
.. autoclass:: einx.nn.haiku.Norm
.. autoclass:: einx.nn.haiku.Dropout
.. autofunction:: einx.nn.haiku.param
Flax
----
.. autofunction:: einx.nn.flax.Linear
.. autofunction:: einx.nn.flax.Norm
.. autofunction:: einx.nn.flax.Dropout
.. autofunction:: einx.nn.flax.param
Torch
-----
.. autoclass:: einx.nn.torch.Linear
.. autoclass:: einx.nn.torch.Norm
.. autoclass:: einx.nn.torch.Dropout
.. autofunction:: einx.nn.torch.param
Equinox
-------
.. autoclass:: einx.nn.equinox.Linear
.. autoclass:: einx.nn.equinox.Norm
.. autoclass:: einx.nn.equinox.Dropout
.. autofunction:: einx.nn.equinox.param
Keras
-----
.. autoclass:: einx.nn.keras.Linear
.. autoclass:: einx.nn.keras.Norm
.. autoclass:: einx.nn.keras.Dropout
.. autofunction:: einx.nn.keras.param
Experimental
============
.. autofunction:: einx.experimental.shard
|
