Index

_

__and__() (pax.experimental.graph.Node method) __call__() (pax.Dropout method) (pax.EMA method) (pax.Embed method) (pax.experimental.graph.GraphModule method) (pax.GroupNorm method) (pax.GRU method) (pax.Identity method) (pax.LayerNorm method) (pax.Linear method) (pax.LSTM method) (pax.MultiHeadAttention method) (pax.Sequential method) (pax.VanillaRNN method) __eq__() (pax.EmptyNode method) (pax.experimental.graph.Node method) __getitem__() (pax.Sequential method) __hash__() (pax.experimental.graph.Node method) __init__() (pax.BatchNorm1D method) (pax.BatchNorm2D method) (pax.Conv1D method) (pax.Conv1DTranspose method) (pax.Conv2D method) (pax.Conv2DTranspose method) (pax.Dropout method) (pax.EMA method) (pax.Embed method) (pax.experimental.Flattener method) (pax.experimental.graph.GraphModule method) (pax.experimental.graph.InputNode method) (pax.GroupNorm method) (pax.GRU method) (pax.LayerNorm method) (pax.Linear method) (pax.LSTM method) (pax.Module method) (pax.MultiHeadAttention method) (pax.RngSeq method) (pax.Sequential method) (pax.VanillaRNN method)

__mod__() (pax.Module method) __or__() (pax.experimental.graph.Node method) (pax.Module method) __rshift__() (pax.experimental.graph.Node method)

A

apply() (pax.Module method) apply_mp_policy() (in module pax)

apply_scaled_gradients() (in module pax.experimental) avg_pool() (in module pax)

B

BatchNorm1D (class in pax) BatchNorm2D (class in pax)

binary_ops() (pax.experimental.graph.Node method) build_graph_module() (in module pax.experimental.graph) build_update_fn() (in module pax)

C

Conv1D (class in pax) Conv1DTranspose (class in pax)

Conv2D (class in pax) Conv2DTranspose (class in pax)

D

default_mp_policy() (in module pax.experimental)

Dropout (class in pax) dtype() (pax.experimental.graph.Node property)

E

EMA (class in pax) Embed (class in pax) EmptyNode (class in pax)

enable_eval_mode() (in module pax) enable_train_mode() (in module pax) eval() (pax.Module method)

F

Flattener (class in pax.experimental)

freeze_parameters() (in module pax)

G

get_or_create() (pax.experimental.LazyModule method) grad() (in module pax)

GraphModule (class in pax.experimental.graph) GroupNorm (class in pax) GRU (class in pax)

I

Identity (class in pax)

initial_state() (pax.GRU method) InputNode (class in pax.experimental.graph)

L

Lambda (class in pax) LayerNorm (class in pax) LazyModule (class in pax.experimental)

Linear (class in pax) load_state_dict() (pax.Module method) load_weights_from_dict() (in module pax.experimental) LSTM (class in pax)

M

max_pool() (in module pax) Module (class in pax)

MultiHeadAttention (class in pax) mutable() (in module pax.experimental)

N

next_rng_key() (in module pax) (pax.RngSeq method)

Node (class in pax.experimental.graph)

P

ParameterModule (class in pax) parameters() (pax.experimental.Flattener method) (pax.Module method) (pax.ParameterModule method)

parameters_method() (in module pax) pure() (in module pax) purecall() (in module pax)

R

replace() (pax.Module method)

replace_node() (pax.Module method) RngSeq (class in pax)

S

save_weights_to_dict() (in module pax.experimental) scan() (in module pax) seed_rng_key() (in module pax) select_parameters() (in module pax) Sequential (class in pax)

set() (pax.Sequential method) shape() (pax.experimental.graph.Node property) state_dict() (pax.Module method) StateModule (class in pax) summary() (pax.Module method)

T

train() (pax.Module method) training() (pax.Module property)

tree_flatten() (pax.EmptyNode method) tree_unflatten() (pax.EmptyNode class method)

U

unfreeze_parameters() (in module pax) unwrap_mp_policy() (in module pax)

update() (pax.experimental.Flattener method) update_parameters() (in module pax) (pax.Module method)

V

value_and_grad() (in module pax)

VanillaRNN (class in pax)