Index _ | A | B | C | D | E | F | G | I | K | L | M | N | O | P | R | S | T | U | V | W _ __call__() (stix.core.coupling.Coupling method) (stix.core.embedded_source_prior.EmbeddedSourcePrior method) (stix.core.loss.Criterion method) (stix.core.loss.CrossEntropyCriterion method) (stix.core.loss.MSECriterion method) (stix.nn.DiTBackbone method) (stix.nn.DiTDecoder method) (stix.nn.DiTEncoder method) (stix.nn.EncoderBackboneDecoderNetwork method) (stix.nn.layers.AdaptiveLayerNorm method) (stix.nn.layers.MultiHeadSelfAttention method) (stix.nn.layers.RegressionHead method) (stix.nn.layers.SinusoidalFourierFeatures method) (stix.nn.layers.SwiGLUFeedForward method) (stix.nn.layers.TransformerBlock method) (stix.nn.MultiHeadDecoder method) (stix.nn.Network method) (stix.nn.SumContextEncoder method) (stix.nn.TimeNoiseContextEncoder method) (stix.sampling.guidance.GuidanceFn method) (stix.sampling.solver.Solver method) (stix.sampling.solver_manual.ManualSolver method) (stix.training.loss_pipeline.LossPipeline method) (stix.training.time_sampler.TimeSampler method) (stix.training.time_sampler.UniformTimeSampler method) A AdaptiveLayerNorm (class in stix.nn.layers) assert_compatible() (stix.core.modality.ModalityRegistry method) assert_fields_set() (stix.core.modality.ModalityRegistry method) assert_num_categories_set_on_discrete_modalities() (stix.core.modality.ModalityRegistry method) attach_logger() (stix.training.training_io_handler.TrainingIOHandler method) attn_mask (stix.typing.data.Batch attribute) B backward_rates (stix.core.generator.TransitionRates attribute) Batch (class in stix.typing.data) best_eval_loss (stix.training.training_loop.TrainingLoop attribute) best_eval_step (stix.training.training_loop.TrainingLoop attribute) BEST_MODEL (stix.training.training_io_handler.LoggingCategory attribute) best_params (stix.training.training_loop.TrainingLoop property) best_step() (stix.training.checkpointer.Checkpointer method) broadcast() (stix.core.modality.ModalityRegistry method) C check_gamma_is_not_zero() (in module stix.core.utils) check_scale_is_non_negative() (in module stix.sampling.utils) checkpoint_dir (stix.training.checkpointer.CheckpointerConfig attribute) (stix.training.training_io_handler.TrainingIOHandlerConfig attribute) Checkpointer (class in stix.training.checkpointer) CheckpointerConfig (class in stix.training.checkpointer) CheckpointRestorationError CLEANUP_AFTER_CKPT_RESTORATION (stix.training.training_io_handler.LoggingCategory attribute) clear_previous_checkpoints (stix.training.training_io_handler.TrainingIOHandlerConfig attribute) Config (stix.training.checkpointer.Checkpointer attribute) (stix.training.training_io_handler.TrainingIOHandler attribute) (stix.training.training_loop.TrainingLoop attribute) context_data (stix.typing.data.Batch attribute) context_dim (stix.nn.NetworkDimsConfig attribute) context_mask (stix.typing.data.Batch attribute) ContinuousBFNOneSidedInterpolant (class in stix.core.interpolant) ContinuousDeterministicInterpolant (class in stix.core.interpolant) ContinuousInterpolant (class in stix.core.interpolant) ContinuousOneSidedStochasticInterpolant (class in stix.core.interpolant) ContinuousStochasticInterpolant (class in stix.core.interpolant) convert_mse_to_rmse_in_logs() (in module stix.training.training_loggers) Coupling (class in stix.core.coupling) Criterion (class in stix.core.loss) CrossEntropyCriterion (class in stix.core.loss) ctmc_euler_step() (in module stix.sampling.steps) D DIFFUSION_FNS (in module stix.sampling.steps) Direction (class in stix.sampling.utils) DiscreteBFNOneSidedInterpolant (class in stix.core.interpolant) DiscreteInterpolant (class in stix.core.interpolant) DiTBackbone (class in stix.nn) DiTDecoder (class in stix.nn) DiTEncoder (class in stix.nn) draw_mixing_index() (stix.core.interpolant.DiscreteInterpolant method) DRIFT_FNS (in module stix.sampling.steps) dtype (stix.nn.NetworkDimsConfig attribute) E ema_decay (stix.training.training_loop_config.TrainingLoopConfig attribute) (stix.training.training_state.TrainingState attribute) ema_model (stix.training.training_loop.TrainingLoop property) ema_params (stix.training.training_state.TrainingState attribute) EmbeddedSourcePrior (class in stix.core.embedded_source_prior) EmbeddedSourceTargetPair (class in stix.typing.variables) EmbeddedVar (in module stix.typing.variables) Embedder (class in stix.core.embedder) embedding_dim (stix.nn.NetworkDimsConfig attribute) enable_async_checkpointing (stix.training.checkpointer.CheckpointerConfig attribute) (stix.training.training_io_handler.TrainingIOHandlerConfig attribute) EncoderBackboneDecoderNetwork (class in stix.nn) eval_every_n_steps (stix.training.training_loop_config.TrainingLoopConfig attribute) EVAL_METRICS (stix.training.training_io_handler.LoggingCategory attribute) F ffn_hidden_dim (stix.nn.NetworkDimsConfig attribute) FlowMatchingOneSidedInterpolant (class in stix.core.interpolant) FlowMatchingTwoSidedInterpolant (class in stix.core.interpolant) FORWARD (stix.sampling.utils.Direction attribute) forward_rates (stix.core.generator.TransitionRates attribute) from_batch() (stix.core.modality.ModalityRegistry class method) from_embeddings_to_logits() (stix.core.embedder.LearnedDiscreteEmbedder method) from_embeddings_to_probs() (stix.core.embedder.LearnedDiscreteEmbedder method) from_embeddings_to_raw() (stix.core.embedder.Embedder method) (stix.core.embedder.IdentityEmbedder method) (stix.core.embedder.LearnedDiscreteEmbedder method) (stix.core.embedder.OneHotDiscreteEmbedder method) from_raw_to_embeddings() (stix.core.embedder.Embedder method) (stix.core.embedder.IdentityEmbedder method) (stix.core.embedder.LearnedDiscreteEmbedder method) (stix.core.embedder.OneHotDiscreteEmbedder method) G GenerativeModel (class in stix.core.gen_model) Generator (class in stix.core.generator) generator_type (stix.core.interpolant.ContinuousDeterministicInterpolant attribute) (stix.core.interpolant.ContinuousStochasticInterpolant attribute) (stix.core.interpolant.DiscreteInterpolant attribute) (stix.core.interpolant.Interpolant attribute) (stix.core.modality.Modality property) get_batch_size() (in module stix.typing.variables) get_classifier_free_guidance_generator() (in module stix.sampling.guidance) get_conditional_rates() (stix.core.interpolant.DiscreteInterpolant method) get_conditional_score() (stix.core.interpolant.ContinuousStochasticInterpolant method) get_conditional_velocity() (stix.core.interpolant.ContinuousInterpolant method) get_conditional_velocity_field() (stix.core.interpolant.ContinuousInterpolant method) get_context_conditional_generator() (in module stix.sampling.guidance) get_debiased_ema() (in module stix.training.ema) get_discrete_classifier_free_guidance_generator() (in module stix.sampling.guidance) get_embeddings() (stix.core.gen_model.GenerativeModel method) get_generator() (stix.core.gen_model.factory.NoiseOneSidedGenerativeModel method) (stix.core.gen_model.factory.PosteriorMixtureGenerativeModel method) (stix.core.gen_model.factory.VelocityAndPosteriorGenerativeModel method) (stix.core.gen_model.factory.VelocityOneSidedGenerativeModel method) (stix.core.gen_model.factory.VelocityTwoSidedGenerativeModel method) (stix.core.gen_model.GenerativeModel method) get_guidance_loss() (stix.core.gen_model.GenerativeModel method) get_intrinsic_guidance_generator() (in module stix.sampling.guidance) get_loss() (stix.core.gen_model.factory.NoiseOneSidedGenerativeModel method) (stix.core.gen_model.factory.PosteriorMixtureGenerativeModel method) (stix.core.gen_model.factory.VelocityAndPosteriorGenerativeModel method) (stix.core.gen_model.factory.VelocityOneSidedGenerativeModel method) (stix.core.gen_model.factory.VelocityTwoSidedGenerativeModel method) (stix.core.gen_model.GenerativeModel method) get_network_output() (stix.core.gen_model.GenerativeModel method) graphdef (stix.training.training_state.TrainingState attribute) GuidanceFn (class in stix.sampling.guidance) I IdentityEmbedder (class in stix.core.embedder) infer_num_samples_from_raw_source() (in module stix.core.utils) integrate() (stix.sampling.solver.Solver method) (stix.sampling.solver_manual.ManualSolver method) integrate_with_trajectory() (stix.sampling.solver_manual.ManualSolver method) Interpolant (class in stix.core.interpolant) interpolant_fn() (stix.core.interpolant.ContinuousInterpolant method) (stix.core.interpolant.ContinuousOneSidedStochasticInterpolant method) (stix.core.interpolant.LinearInterpolant method) interpolate() (stix.core.interpolant.ContinuousInterpolant method) (stix.core.interpolant.DiscreteInterpolant method) (stix.core.interpolant.Interpolant method) is_discrete (stix.typing.data.Batch attribute) is_one_hot_vector() (in module stix.core.utils) K keep_best (stix.training.checkpointer.CheckpointerConfig attribute) (stix.training.training_io_handler.TrainingIOHandlerConfig attribute) L latest_step() (stix.training.checkpointer.Checkpointer method) LearnedDiscreteEmbedder (class in stix.core.embedder) LinearDeterministicInterpolant (class in stix.core.interpolant) LinearInterpolant (class in stix.core.interpolant) LinearStochasticInterpolant (class in stix.core.interpolant) log() (stix.training.training_io_handler.TrainingIOHandler method) log_metrics_to_line() (in module stix.training.training_loggers) log_metrics_to_table() (in module stix.training.training_loggers) LoggingCategory (class in stix.training.training_io_handler) loss_mask (stix.typing.data.Batch attribute) LossPipeline (class in stix.training.loss_pipeline) M ManualSolver (class in stix.sampling.solver_manual) ManualSolverConfig (class in stix.sampling.solver_manual) map() (stix.core.modality.ModalityRegistry method) Mask (in module stix.typing.core) mask_index (stix.core.interpolant.MaskDiscreteInterpolant property) MaskDiscreteInterpolant (class in stix.core.interpolant) max_to_keep (stix.training.checkpointer.CheckpointerConfig attribute) (stix.training.training_io_handler.TrainingIOHandlerConfig attribute) mixture_conditionals() (stix.core.interpolant.DiscreteInterpolant method) (stix.core.interpolant.MaskDiscreteInterpolant method) (stix.core.interpolant.UniformDiscreteInterpolant method) mixture_distributions_from_target_posterior() (stix.core.interpolant.MaskDiscreteInterpolant method) (stix.core.interpolant.UniformDiscreteInterpolant method) Modality (class in stix.core.modality) ModalityRegistry (class in stix.core.modality) model_config (stix.nn.NetworkDimsConfig attribute) (stix.training.checkpointer.CheckpointerConfig attribute) (stix.training.training_io_handler.TrainingIOHandlerConfig attribute) (stix.training.training_loop_config.TrainingLoopConfig attribute) module stix stix.sampling.guidance stix.sampling.steps stix.typing.core MSECriterion (class in stix.core.loss) MultiHeadDecoder (class in stix.nn) MultiHeadSelfAttention (class in stix.nn.layers) N Network (class in stix.nn) NetworkDimsConfig (class in stix.nn) noise_from_score() (stix.core.interpolant.LinearStochasticInterpolant method) noise_from_target() (stix.core.interpolant.OneSidedLinearStochasticInterpolant method) noise_from_velocity() (stix.core.interpolant.OneSidedLinearStochasticInterpolant method) noise_from_velocity_field() (stix.core.interpolant.OneSidedLinearStochasticInterpolant method) NoiseOneSidedGenerativeModel (class in stix.core.gen_model.factory) NoiseVar (in module stix.typing.variables) num_gradient_accumulation_steps (stix.training.training_loop_config.TrainingLoopConfig attribute) num_steps (stix.training.training_loop_config.TrainingLoopConfig attribute) (stix.training.training_state.TrainingState attribute) O OneHotDiscreteEmbedder (class in stix.core.embedder) OneSidedInterpolant (class in stix.core.interpolant) OneSidedLinearStochasticInterpolant (class in stix.core.interpolant) opt_state (stix.training.training_state.TrainingState attribute) P params (stix.training.training_state.TrainingState attribute) PosteriorMixtureGenerativeModel (class in stix.core.gen_model.factory) PRNGKeyArray (in module stix.typing.core) R random_seed (stix.training.training_loop_config.TrainingLoopConfig attribute) rates_from_mixture_distributions() (stix.core.interpolant.DiscreteInterpolant method) rates_from_target_posterior() (stix.core.interpolant.MaskDiscreteInterpolant method) (stix.core.interpolant.UniformDiscreteInterpolant method) raw_batch (stix.typing.data.Batch attribute) RawCtsVar (in module stix.typing.variables) RawDisVar (in module stix.typing.variables) RawSourceTargetPair (class in stix.typing.variables) RawVar (in module stix.typing.variables) read_run_metadata() (in module stix.training.checkpointer) reduce_modality_losses() (in module stix.core.loss) RegressionHead (class in stix.nn.layers) replace() (stix.core.generator.Generator method) (stix.core.generator.TransitionRates method) (stix.core.generator.Velocity method) (stix.core.generator.VelocityAndScore method) (stix.training.training_state.TrainingState method) (stix.typing.data.Batch method) restore() (stix.training.checkpointer.Checkpointer method) restore_checkpoint() (stix.training.training_io_handler.TrainingIOHandler method) restore_checkpoint_if_exists (stix.training.training_io_handler.TrainingIOHandlerConfig attribute) restore_dir (stix.training.training_io_handler.TrainingIOHandlerConfig attribute) restore_ema() (stix.training.checkpointer.Checkpointer method) restore_model() (stix.training.training_loop.TrainingLoop method) restore_optimizer_state (stix.training.training_io_handler.TrainingIOHandlerConfig attribute) REVERSE (stix.sampling.utils.Direction attribute) rng_key (stix.training.training_state.TrainingState attribute) run() (stix.training.training_loop.TrainingLoop method) run_eval_at_start (stix.training.training_loop_config.TrainingLoopConfig attribute) S sample_initial_state() (stix.core.interpolant.Interpolant method) (stix.core.interpolant.LinearInterpolant method) (stix.core.interpolant.MaskDiscreteInterpolant method) (stix.core.interpolant.UniformDiscreteInterpolant method) (stix.core.modality.ModalityRegistry method) sample_noise() (stix.core.interpolant.ContinuousInterpolant method) (stix.core.interpolant.DiscreteInterpolant method) (stix.core.interpolant.Interpolant method) save() (stix.training.checkpointer.Checkpointer method) save_checkpoint() (stix.training.training_io_handler.TrainingIOHandler method) save_interval_steps (stix.training.checkpointer.CheckpointerConfig attribute) (stix.training.training_io_handler.TrainingIOHandlerConfig attribute) save_run_metadata() (in module stix.training.checkpointer) Scalar (in module stix.typing.core) score (stix.core.generator.VelocityAndScore attribute) score_from_noise() (stix.core.interpolant.LinearStochasticInterpolant method) score_from_target() (stix.core.interpolant.OneSidedLinearStochasticInterpolant method) score_from_velocity() (stix.core.interpolant.OneSidedLinearStochasticInterpolant method) score_from_velocity_field() (stix.core.interpolant.OneSidedLinearStochasticInterpolant method) set() (stix.core.modality.ModalityRegistry method) Shape (in module stix.typing.core) SinusoidalFourierFeatures (class in stix.nn.layers) Solver (class in stix.sampling.solver) SolverConfig (class in stix.sampling.solver) source (stix.typing.variables.EmbeddedSourceTargetPair attribute) (stix.typing.variables.RawSourceTargetPair attribute) split_and_project_key() (stix.core.modality.ModalityRegistry method) STEP_FNS (in module stix.sampling.steps) step_to_restore (stix.training.training_io_handler.TrainingIOHandlerConfig attribute) stix module stix.sampling.guidance module stix.sampling.steps module stix.typing.core module StochasticFlowMatchingTwoSidedInterpolant (class in stix.core.interpolant) SumContextEncoder (class in stix.nn) SwiGLUFeedForward (class in stix.nn.layers) SYSTEM_METRICS (stix.training.training_io_handler.LoggingCategory attribute) T target (stix.typing.variables.EmbeddedSourceTargetPair attribute) (stix.typing.variables.RawSourceTargetPair attribute) target_from_noise() (stix.core.interpolant.OneSidedLinearStochasticInterpolant method) target_from_score() (stix.core.interpolant.OneSidedLinearStochasticInterpolant method) target_from_velocity() (stix.core.interpolant.OneSidedLinearStochasticInterpolant method) target_from_velocity_field() (stix.core.interpolant.OneSidedLinearStochasticInterpolant method) TEST_METRICS (stix.training.training_io_handler.LoggingCategory attribute) Time (in module stix.typing.core) TimeNoiseContextEncoder (class in stix.nn) TimeSampler (class in stix.training.time_sampler) TRAIN_METRICS (stix.training.training_io_handler.LoggingCategory attribute) training_state (stix.training.training_loop.TrainingLoop attribute) TrainingIOHandler (class in stix.training.training_io_handler) TrainingIOHandlerConfig (class in stix.training.training_io_handler) TrainingLoop (class in stix.training.training_loop) TrainingLoopConfig (class in stix.training.training_loop_config) TrainingState (class in stix.training.training_state) TransformerBlock (class in stix.nn.layers) TransitionRates (class in stix.core.generator) treedef (stix.core.modality.ModalityRegistry property) U UniformDiscreteInterpolant (class in stix.core.interpolant) UniformTimeSampler (class in stix.training.time_sampler) use_ema_params_for_eval (stix.training.training_loop_config.TrainingLoopConfig attribute) V validate_generator_types() (in module stix.sampling.utils) Var (in module stix.typing.core) VarianceExplodingDiffusionOneSidedInterpolant (class in stix.core.interpolant) VarType (class in stix.typing.variable_type) Velocity (class in stix.core.generator) velocity (stix.core.generator.Velocity attribute) (stix.core.generator.VelocityAndScore attribute) velocity_and_score_diffusion() (in module stix.sampling.steps) velocity_and_score_drift() (in module stix.sampling.steps) velocity_and_score_euler_maruyama_step() (in module stix.sampling.steps) velocity_diffusion() (in module stix.sampling.steps) velocity_drift() (in module stix.sampling.steps) velocity_euler_step() (in module stix.sampling.steps) velocity_field_from_noise() (stix.core.interpolant.OneSidedLinearStochasticInterpolant method) velocity_field_from_score() (stix.core.interpolant.OneSidedLinearStochasticInterpolant method) velocity_field_from_target() (stix.core.interpolant.OneSidedLinearStochasticInterpolant method) velocity_field_from_velocity() (stix.core.interpolant.LinearDeterministicInterpolant method) (stix.core.interpolant.OneSidedLinearStochasticInterpolant method) velocity_field_from_velocity_and_noise() (stix.core.interpolant.LinearStochasticInterpolant method) velocity_field_from_velocity_and_score() (stix.core.interpolant.LinearStochasticInterpolant method) velocity_from_noise() (stix.core.interpolant.OneSidedLinearStochasticInterpolant method) velocity_from_score() (stix.core.interpolant.OneSidedLinearStochasticInterpolant method) velocity_from_target() (stix.core.interpolant.OneSidedLinearStochasticInterpolant method) velocity_from_velocity_field() (stix.core.interpolant.LinearDeterministicInterpolant method) (stix.core.interpolant.OneSidedLinearStochasticInterpolant method) velocity_from_velocity_field_and_noise() (stix.core.interpolant.LinearStochasticInterpolant method) velocity_from_velocity_field_and_score() (stix.core.interpolant.LinearStochasticInterpolant method) VelocityAndPosteriorGenerativeModel (class in stix.core.gen_model.factory) VelocityAndScore (class in stix.core.generator) VelocityOneSidedGenerativeModel (class in stix.core.gen_model.factory) VelocityTwoSidedGenerativeModel (class in stix.core.gen_model.factory) W wait_until_finished() (stix.training.checkpointer.Checkpointer method) (stix.training.training_io_handler.TrainingIOHandler method)