Index

_

__call__() (haiku.AvgPool method) (haiku.BatchApply method) (haiku.BatchNorm method) (haiku.Bias method) (haiku.ConvND method) (haiku.ConvNDTranspose method) (haiku.Deferred method) (haiku.EMAParamsTree method) (haiku.Embed method) (haiku.ExponentialMovingAverage method) (haiku.GroupNorm method) (haiku.GRU method) (haiku.IdentityCore method) (haiku.initializers.Constant method) (haiku.initializers.Identity method) (haiku.initializers.Orthogonal method) (haiku.initializers.RandomNormal method) (haiku.initializers.RandomUniform method) (haiku.initializers.TruncatedNormal method) (haiku.initializers.UniformScaling method) (haiku.initializers.VarianceScaling method) (haiku.LayerNorm method) (haiku.Linear method) (haiku.LSTM method) (haiku.MaxPool method) (haiku.MultiHeadAttention method) (haiku.nets.MLP method) (haiku.nets.MobileNetV1 method) (haiku.nets.ResNet method) (haiku.nets.ResNet.BlockGroup method) (haiku.nets.ResNet.BlockV1 method) (haiku.nets.ResNet.BlockV2 method) (haiku.nets.VectorQuantizer method) (haiku.nets.VectorQuantizerEMA method) (haiku.ResetCore method) (haiku.Reshape method) (haiku.RMSNorm method) (haiku.RNNCore method) (haiku.SeparableDepthwiseConv2D method) (haiku.Sequential method) (haiku.SNParamsTree method) (haiku.SpectralNorm method) (haiku.VanillaRNN method) __delattr__() (haiku.Deferred method) (haiku.experimental.ArraySpec method) (haiku.experimental.MethodInvocation method) (haiku.experimental.ModuleDetails method) __eq__() (haiku.experimental.ArraySpec method) (haiku.experimental.MethodInvocation method) (haiku.experimental.ModuleDetails method) __hash__() (haiku.experimental.ArraySpec method) (haiku.experimental.MethodInvocation method) (haiku.experimental.ModuleDetails method) __init__() (haiku.AvgPool method) (haiku.BatchApply method) (haiku.BatchNorm method) (haiku.Bias method) (haiku.Conv1D method) (haiku.Conv1DLSTM method) (haiku.Conv1DTranspose method) (haiku.Conv2D method) (haiku.Conv2DLSTM method) (haiku.Conv2DTranspose method) (haiku.Conv3D method) (haiku.Conv3DLSTM method) (haiku.Conv3DTranspose method) (haiku.ConvND method) (haiku.ConvNDTranspose method) (haiku.DeepRNN method) (haiku.Deferred method) (haiku.DepthwiseConv1D method) (haiku.DepthwiseConv2D method) (haiku.DepthwiseConv3D method) (haiku.EMAParamsTree method) (haiku.Embed method) (haiku.experimental.ArraySpec method) (haiku.experimental.MethodInvocation method) (haiku.experimental.ModuleDetails method) (haiku.ExponentialMovingAverage method) (haiku.Flatten method) (haiku.GroupNorm method) (haiku.GRU method) (haiku.initializers.Constant method) (haiku.initializers.Identity method) (haiku.initializers.Orthogonal method) (haiku.initializers.RandomNormal method) (haiku.initializers.RandomUniform method) (haiku.initializers.TruncatedNormal method) (haiku.initializers.UniformScaling method) (haiku.initializers.VarianceScaling method) (haiku.InstanceNorm method) (haiku.LayerNorm method) (haiku.Linear method) (haiku.LSTM method) (haiku.MaxPool method) (haiku.Module method) (haiku.MultiHeadAttention method) (haiku.nets.MLP method) (haiku.nets.MobileNetV1 method) (haiku.nets.ResNet method) (haiku.nets.ResNet.BlockGroup method) (haiku.nets.ResNet.BlockV1 method) (haiku.nets.ResNet.BlockV2 method) (haiku.nets.ResNet101 method) (haiku.nets.ResNet152 method) (haiku.nets.ResNet18 method) (haiku.nets.ResNet200 method) (haiku.nets.ResNet34 method) (haiku.nets.ResNet50 method) (haiku.nets.VectorQuantizer method) (haiku.nets.VectorQuantizerEMA method) (haiku.PRNGSequence method) (haiku.ResetCore method) (haiku.Reshape method) (haiku.RMSNorm method) (haiku.SeparableDepthwiseConv2D method) (haiku.Sequential method) (haiku.SNParamsTree method) (haiku.SpectralNorm method) (haiku.VanillaRNN method)

__next__() (haiku.PRNGSequence method) __post_init__() (haiku.Module method) __setattr__() (haiku.Deferred method) (haiku.experimental.ArraySpec method) (haiku.experimental.MethodInvocation method) (haiku.experimental.ModuleDetails method)

A

abstract_to_dot() (in module haiku.experimental) apply (haiku.MultiTransformed attribute) (haiku.MultiTransformedWithState attribute) (haiku.Transformed attribute) (haiku.TransformedWithState attribute) args_spec (haiku.experimental.MethodInvocation attribute)

ARRAY_INDEX (haiku.EmbedLookupStyle attribute) ArraySpec (class in haiku.experimental) as_html() (in module haiku.experimental.jaxpr_info) as_html_page() (in module haiku.experimental.jaxpr_info) avg_pool() (in module haiku) AvgPool (class in haiku)

B

BatchApply (class in haiku)

BatchNorm (class in haiku) Bias (class in haiku)

C

call_stack (haiku.experimental.MethodInvocation attribute) causal() (in module haiku.pad) cell (haiku.LSTMState attribute) check_jax_usage() (in module haiku.experimental) clear_policy() (in module haiku.mixed_precision) commitment_cost (haiku.nets.VectorQuantizer attribute) (haiku.nets.VectorQuantizerEMA attribute) cond() (in module haiku) Constant (class in haiku.initializers) context (haiku.experimental.MethodInvocation attribute) context() (in module haiku.config) Conv1D (class in haiku) Conv1DLSTM (class in haiku) Conv1DTranspose (class in haiku) Conv2D (class in haiku)

Conv2DLSTM (class in haiku) Conv2DTranspose (class in haiku) Conv3D (class in haiku) Conv3DLSTM (class in haiku) Conv3DTranspose (class in haiku) ConvND (class in haiku) ConvNDTranspose (class in haiku) create() (in module haiku.pad) create_from_padfn() (in module haiku.pad) create_from_tuple() (in module haiku.pad) css() (in module haiku.experimental.jaxpr_info) current_name() (in module haiku) current_policy() (in module haiku.mixed_precision) custom_creator() (in module haiku) custom_getter() (in module haiku) custom_setter() (in module haiku)

D

decay (haiku.nets.VectorQuantizerEMA attribute) deep_rnn_with_skip_connections() (in module haiku) DeepRNN (class in haiku) Deferred (class in haiku) DepthwiseConv1D (class in haiku)

DepthwiseConv2D (class in haiku) DepthwiseConv3D (class in haiku) DO_NOT_STORE (in module haiku) dropout() (in module haiku) dtype (haiku.experimental.ArraySpec attribute) dynamic_unroll() (in module haiku)

E

EMAParamsTree (class in haiku) Embed (class in haiku) embedding_dim (haiku.nets.VectorQuantizer attribute) (haiku.nets.VectorQuantizerEMA attribute) EmbedLookupStyle (class in haiku)

epsilon (haiku.nets.VectorQuantizerEMA attribute) eval_shape() (in module haiku) eval_summary() (in module haiku.experimental) expand_apply() (in module haiku) ExponentialMovingAverage (class in haiku) Expression (class in haiku.experimental.jaxpr_info)

F

fast_eval_shape() (in module haiku.experimental) filter() (in module haiku.data_structures) Flatten (class in haiku) flatten_flax_to_haiku() (in module haiku.experimental.flax) force_name() (in module haiku)

fori_loop() (in module haiku) format_module() (in module haiku.experimental.jaxpr_info) full() (in module haiku.pad) full_name (haiku.GetterContext attribute) (haiku.SetterContext attribute)

G

get_channel_index() (in module haiku) get_current_state() (in module haiku) get_initial_state() (in module haiku) get_parameter() (in module haiku) get_params() (in module haiku)

get_policy() (in module haiku.mixed_precision) get_state() (in module haiku) GetterContext (class in haiku) grad() (in module haiku) GroupNorm (class in haiku) GRU (class in haiku)

H

haiku module, [1], [2] haiku.config module haiku.data_structures module haiku.experimental module haiku.experimental.flax module haiku.experimental.jaxpr_info module

haiku.initializers module haiku.mixed_precision module haiku.nets module haiku.pad module haiku.testing module hidden (haiku.LSTMState attribute)

I

Identity (class in haiku.initializers) IdentityCore (class in haiku) init (haiku.MultiTransformed attribute) (haiku.MultiTransformedWithState attribute) (haiku.Transformed attribute) (haiku.TransformedWithState attribute) initial_state() (haiku.GRU method) (haiku.IdentityCore method) (haiku.LSTM method) (haiku.ResetCore method) (haiku.RNNCore method) (haiku.VanillaRNN method)

initialize() (haiku.ExponentialMovingAverage method) Initializer (in module haiku.initializers) InstanceNorm (class in haiku) intercept_methods() (in module haiku) is_padfn() (in module haiku.pad) is_subset() (in module haiku.data_structures)

J

js() (in module haiku.experimental.jaxpr_info)

K

kwargs_spec (haiku.experimental.MethodInvocation attribute)

L

layer_stack (class in haiku) LayerNorm (class in haiku) LayerStackTransparencyMapping (class in haiku) lift() (in module haiku) (in module haiku.experimental.flax) lift_with_state() (in module haiku)

lifted_prefix_name (haiku.GetterContext attribute) (haiku.SetterContext attribute) LiftWithStateUpdater (class in haiku) Linear (class in haiku) LSTM (class in haiku) LSTMState (class in haiku)

M

make_model_info() (in module haiku.experimental.jaxpr_info) map() (in module haiku) (in module haiku.data_structures) max_pool() (in module haiku) MaxPool (class in haiku) maybe_get_rng_sequence_state() (in module haiku) maybe_next_rng_key() (in module haiku) merge() (in module haiku.data_structures) method_name (haiku.experimental.ModuleDetails attribute) (haiku.MethodContext attribute) MethodContext (class in haiku) MethodInvocation (class in haiku.experimental) MLP (class in haiku.nets) MobileNetV1 (class in haiku.nets) module haiku, [1], [2] haiku.config haiku.data_structures haiku.experimental haiku.experimental.flax haiku.experimental.jaxpr_info haiku.initializers haiku.mixed_precision haiku.nets haiku.pad haiku.testing

Module (class in haiku) (class in haiku.experimental.flax) (class in haiku.experimental.jaxpr_info) module (haiku.experimental.ModuleDetails attribute) (haiku.GetterContext attribute) (haiku.MethodContext attribute) (haiku.SetterContext attribute) module_auto_repr() (in module haiku.experimental) module_details (haiku.experimental.MethodInvocation attribute) module_name (haiku.GetterContext attribute) (haiku.SetterContext attribute) ModuleDetails (class in haiku.experimental) ModuleProtocol (class in haiku) multi_transform() (in module haiku) multi_transform_with_state() (in module haiku) MultiHeadAttention (class in haiku) multinomial() (in module haiku) MultiTransformed (class in haiku) MultiTransformedWithState (class in haiku) MutableParams (in module haiku) MutableState (in module haiku)

N

name (haiku.GetterContext attribute) (haiku.SetterContext attribute) name_like() (in module haiku) name_scope() (in module haiku)

next() (haiku.PRNGSequence method) next_rng_key() (in module haiku) next_rng_keys() (in module haiku) num_embeddings (haiku.nets.VectorQuantizer attribute) (haiku.nets.VectorQuantizerEMA attribute)

O

ONE_HOT (haiku.EmbedLookupStyle attribute) one_hot() (in module haiku) optimize_rng_use() (in module haiku.experimental) orig_class (haiku.MethodContext attribute) orig_method (haiku.MethodContext attribute) original_dtype (haiku.GetterContext attribute) (haiku.SetterContext attribute)

original_init (haiku.GetterContext attribute) original_shape (haiku.GetterContext attribute) (haiku.SetterContext attribute) Orthogonal (class in haiku.initializers) output_spec (haiku.experimental.MethodInvocation attribute)

P

PadFn (in module haiku.pad) params (haiku.experimental.ModuleDetails attribute) Params (in module haiku) params_dict() (haiku.Module method)

partition() (in module haiku.data_structures) partition_n() (in module haiku.data_structures) PRNGSequence (class in haiku) push_policy() (in module haiku.mixed_precision)

Q

quantize() (haiku.nets.VectorQuantizer method) (haiku.nets.VectorQuantizerEMA method)

R

RandomNormal (class in haiku.initializers) RandomUniform (class in haiku.initializers) remat() (in module haiku) replace_rng_sequence_state() (in module haiku) reserve() (haiku.PRNGSequence method) reserve_rng_keys() (in module haiku) ResetCore (class in haiku) Reshape (class in haiku) ResNet (class in haiku.nets) ResNet.BlockGroup (class in haiku.nets) ResNet.BlockV1 (class in haiku.nets) ResNet.BlockV2 (class in haiku.nets)

ResNet101 (class in haiku.nets) ResNet152 (class in haiku.nets) ResNet18 (class in haiku.nets) ResNet200 (class in haiku.nets) ResNet34 (class in haiku.nets) ResNet50 (class in haiku.nets) reverse() (haiku.nets.MLP method) reverse_causal() (in module haiku.pad) RMSNorm (class in haiku) rng_reserve_size() (in module haiku.experimental) RNNCore (class in haiku) running_init() (in module haiku)

S

same() (in module haiku.pad) scan() (in module haiku) SeparableDepthwiseConv2D (class in haiku) Sequential (class in haiku) set() (in module haiku.config) set_policy() (in module haiku.mixed_precision) set_state() (in module haiku) SetterContext (class in haiku)

shape (haiku.experimental.ArraySpec attribute) SNParamsTree (class in haiku) SpectralNorm (class in haiku) state (haiku.experimental.ModuleDetails attribute) State (in module haiku) state_dict() (haiku.Module method) static_unroll() (in module haiku) SupportsCall (class in haiku) switch() (in module haiku)

T

tabulate() (in module haiku.experimental) target (haiku.Deferred property) to_dot() (in module haiku) to_haiku_dict() (in module haiku.data_structures) to_immutable_dict() (in module haiku.data_structures) to_module() (in module haiku) to_mutable_dict() (in module haiku.data_structures) transform() (in module haiku) transform_and_run() (in module haiku.testing)

transform_with_state() (in module haiku) Transformed (class in haiku) TransformedWithState (class in haiku) transparent() (in module haiku) transparent_lift() (in module haiku) transparent_lift_with_state() (in module haiku) traverse() (in module haiku.data_structures) tree_bytes() (in module haiku.data_structures) tree_size() (in module haiku.data_structures) TruncatedNormal (class in haiku.initializers)

U

UniformScaling (class in haiku.initializers)

V

valid() (in module haiku.pad) value_and_grad() (in module haiku) VanillaRNN (class in haiku)

VarianceScaling (class in haiku.initializers) VectorQuantizer (class in haiku.nets) VectorQuantizerEMA (class in haiku.nets) vmap() (in module haiku)

W

while_loop() (in module haiku) with_empty_state() (in module haiku)

with_rng() (in module haiku) without_apply_rng() (in module haiku) without_state() (in module haiku)