Index

A

apply() (vit_flax.layers.FeedForward method) (vit_flax.layers.PreNorm method) (vit_flax.layers.Residual method) (vit_flax.layers.Transformer method) (vit_flax.vit.ViT method)

C

call() (vit_flax.layers.FeedForward class method) (vit_flax.layers.PreNorm class method) (vit_flax.layers.Residual class method) (vit_flax.layers.Transformer class method) (vit_flax.vit.ViT class method) create() (vit_flax.layers.FeedForward class method) (vit_flax.layers.PreNorm class method) (vit_flax.layers.Residual class method) (vit_flax.layers.Transformer class method) (vit_flax.vit.ViT class method)

create_by_shape() (vit_flax.layers.FeedForward class method) (vit_flax.layers.PreNorm class method) (vit_flax.layers.Residual class method) (vit_flax.layers.Transformer class method) (vit_flax.vit.ViT class method)

F

FeedForward (class in vit_flax.layers)

I

init() (vit_flax.layers.FeedForward class method) (vit_flax.layers.PreNorm class method) (vit_flax.layers.Residual class method) (vit_flax.layers.Transformer class method) (vit_flax.vit.ViT class method)

init_by_shape() (vit_flax.layers.FeedForward class method) (vit_flax.layers.PreNorm class method) (vit_flax.layers.Residual class method) (vit_flax.layers.Transformer class method) (vit_flax.vit.ViT class method)

P

partial() (vit_flax.layers.FeedForward class method) (vit_flax.layers.PreNorm class method) (vit_flax.layers.Residual class method) (vit_flax.layers.Transformer class method) (vit_flax.vit.ViT class method)

PreNorm (class in vit_flax.layers)

R

Residual (class in vit_flax.layers)

T

Transformer (class in vit_flax.layers)

V

ViT (class in vit_flax.vit) vit_flax (module)

vit_flax.layers (module) vit_flax.vit (module)