Index

A

apply() (mlax.module.Module method)

apply_attention_weights() (in module mlax.nn.functional) avg_pool() (in module mlax.nn.functional)

B

Bias (class in mlax.nn.bias)

C

combine() (mlax.module.Module method)

Conv (class in mlax.nn.conv)

D

dot_product_attention_logits() (in module mlax.nn.functional)

dropout() (in module mlax.nn.functional)

E

Embed (class in mlax.nn.embed)

F

F (class in mlax.nn.f) filter() (mlax.module.Module method) filter_with_path() (mlax.module.Module method) forward() (mlax.module.Module method) (mlax.nn.bias.Bias method) (mlax.nn.conv.Conv method) (mlax.nn.embed.Embed method) (mlax.nn.f.F method) (mlax.nn.f.FRng method) (mlax.nn.linear.Linear method) (mlax.nn.parallel.Parallel method) (mlax.nn.parallel.ParallelRng method) (mlax.nn.scaler.Scaler method) (mlax.nn.series.Series method) (mlax.nn.series.SeriesRng method) (mlax.nn.z_norm.ZNorm method)

FRng (class in mlax.nn.f)

I

identity() (in module mlax.nn.functional) is_leaf_param() (in module mlax.module)

is_non_trainable_param() (in module mlax.module) is_trainable_param() (in module mlax.module)

L

Linear (class in mlax.nn.linear)

M

max_pool() (in module mlax.nn.functional) mlax module mlax.module module mlax.nn module mlax.nn.bias module mlax.nn.conv module mlax.nn.embed module mlax.nn.f module mlax.nn.functional module mlax.nn.linear module mlax.nn.parallel module

mlax.nn.scaler module mlax.nn.series module mlax.nn.z_norm module module mlax mlax.module mlax.nn mlax.nn.bias mlax.nn.conv mlax.nn.embed mlax.nn.f mlax.nn.functional mlax.nn.linear mlax.nn.parallel mlax.nn.scaler mlax.nn.series mlax.nn.z_norm Module (class in mlax.module)

P

Parallel (class in mlax.nn.parallel) ParallelRng (class in mlax.nn.parallel) Parameter (class in mlax.module)

partition() (mlax.module.Module method) partition_with_path() (mlax.module.Module method) pool() (in module mlax.nn.functional)

S

Scaler (class in mlax.nn.scaler) Series (class in mlax.nn.series) SeriesRng (class in mlax.nn.series) setup() (mlax.module.Module method) (mlax.nn.bias.Bias method) (mlax.nn.conv.Conv method) (mlax.nn.embed.Embed method) (mlax.nn.f.F method) (mlax.nn.f.FRng method) (mlax.nn.linear.Linear method) (mlax.nn.parallel.Parallel method) (mlax.nn.parallel.ParallelRng method) (mlax.nn.scaler.Scaler method) (mlax.nn.series.Series method) (mlax.nn.series.SeriesRng method) (mlax.nn.z_norm.ZNorm method)

sum_pool() (in module mlax.nn.functional)

T

tree_flatten() (mlax.module.Parameter method) tree_flatten_with_keys() (mlax.module.Module method)

tree_unflatten() (mlax.module.Module class method) (mlax.module.Parameter class method)

Z

z_norm() (in module mlax.nn.functional)

ZNorm (class in mlax.nn.z_norm)