Core Components

This section contains documentation for core components.

Modules

Benchmarks

• base
• runner
• types

Cli

• config_commands

Configuration

• environment
• gan
• migrate_configs
• unified
• validation

Core

• adapters
• base
• checkpointing
• device_manager
• device_testing
• interfaces
• jax_config
• logging
• parallelism
• performance
• types

Distributions

• base
• continuous
• discrete
• mixture
• transformations

General

• precision_recall

Image

• fid
• inception_score

Layers

• causal
• flash_attention
• positional
• residual
• resnet
• transformers

Losses

• adversarial
• base
• composable
• divergence
• geometric
• perceptual
• reconstruction
• regularization

Management

• templates
• versioning

Metrics

• base
• distance
• information
• metric_ops
• pipeline
• quality
• registry
• statistical

Protocols

• benchmarks
• configuration
• evaluation
• metrics

Sampling

• ancestral
• base
• blackjax_samplers
• diffusion
• mcmc
• ode
• sde

Text

• perplexity

Statistics

• Total modules: 79
• Total classes: 131
• Total functions: 543