Skip to main content
Back to top
Ctrl
+
K
kauldron
kauldron
Guides
Introduction
Train, eval, randomness
Sharding & Model parallelism
Checkpoint
Data pipelines
Konfig principles
Modules
Konfig - Demo
kontext
PyGrain data pipeline
Metrics
Klinen - Torch-like API for flax
OO Random API
Type Annotations and Shape checking
Links
GitHub
Issues
API
kd
kd._filter_logs
kd._filter_logs.add_filter
kd.ckpts
kd.ckpts.AbstractPartialLoader
kd.ckpts.Checkpointer
kd.ckpts.MultiTransform
kd.ckpts.NoopCheckpointer
kd.ckpts.NoopTransform
kd.ckpts.PartialKauldronLoader
kd.ckpts.items
kd.ckpts.items.CheckpointItem
kd.ckpts.items.StandardCheckpointItem
kd.ckpts.items.TopLevelCheckpointItem
kd.ckpts.workdir_from_xid
kd.contrib
kd.data
kd.data.AddConstants
kd.data.BatchSize
kd.data.Cast
kd.data.ElementWiseTransform
kd.data.Elements
kd.data.FilterTransform
kd.data.Gather
kd.data.InMemoryPipeline
kd.data.IterableDataset
kd.data.MapTransform
kd.data.Pipeline
kd.data.Rearrange
kd.data.Resize
kd.data.TreeFlattenWithPath
kd.data.ValueRange
kd.data.iters
kd.data.iters.Iterator
kd.data.iters.NonCheckpointableIterator
kd.data.iters.PyGrainIterator
kd.data.iters.TFDataIterator
kd.data.py
kd.data.py.AddConstants
kd.data.py.Cast
kd.data.py.DataSource
kd.data.py.DataSourceBase
kd.data.py.ElementWiseTransform
kd.data.py.Elements
kd.data.py.Gather
kd.data.py.HuggingFace
kd.data.py.Json
kd.data.py.Mix
kd.data.py.PyGrainPipeline
kd.data.py.Rearrange
kd.data.py.Resize
kd.data.py.SliceDataset
kd.data.py.Tfds
kd.data.py.TreeFlattenWithPath
kd.data.py.ValueRange
kd.evals
kd.evals.CollectionKeys
kd.evals.Evaluator
kd.evals.EvaluatorBase
kd.evals.EveryNSteps
kd.evals.FewShotEvaluator
kd.evals.Once
kd.evals.RunStrategy
kd.evals.StandaloneEveryCheckpoint
kd.evals.StandaloneLastCheckpoint
kd.from_xid
kd.from_xid.get_cfg
kd.from_xid.get_element_spec
kd.from_xid.get_resolved
kd.from_xid.get_workdir
kd.inspect
kd.inspect.Profiler
kd.inspect.get_batch_stats
kd.inspect.get_colab_model_overview
kd.inspect.get_connection_graph
kd.inspect.json_spec_like
kd.inspect.lower_trainstep
kd.inspect.plot_batch
kd.inspect.plot_context
kd.inspect.plot_schedules
kd.inspect.plot_sharding
kd.inspect.show_trainer_info
kd.kdash
kd.kdash.BuildContext
kd.kdash.DashboardsBase
kd.kdash.MetricDashboards
kd.kdash.MultiDashboards
kd.kdash.NoopDashboard
kd.kdash.Plot
kd.kdash.SingleDashboard
kd.kdash.build_and_upload
kd.knn
kd.knn.Dense
kd.knn.Dropout
kd.knn.Intermediate
kd.knn.Module
kd.knn.Sequential
kd.knn.convert
kd.konfig
kd.konfig.ConfigDict
kd.konfig.ConfigDictLike
kd.konfig.DEFINE_config_file
kd.konfig.WithRef
kd.konfig._default_values
kd.konfig.imports
kd.konfig.mock_modules
kd.konfig.placeholder
kd.konfig.ref_copy
kd.konfig.ref_fn
kd.konfig.register_aliases
kd.konfig.register_default_values
kd.konfig.required
kd.konfig.resolve
kd.konfig.set_lazy_imported_modules
kd.kontext
kd.kontext.Context
kd.kontext.GlobPath
kd.kontext.Key
kd.kontext.KeyTree
kd.kontext.Path
kd.kontext.REQUIRED
kd.kontext.filter_by_path
kd.kontext.flatten_with_path
kd.kontext.get_by_path
kd.kontext.get_keypaths
kd.kontext.is_key_annotated
kd.kontext.path_builder_from
kd.kontext.resolve_from_keyed_obj
kd.kontext.resolve_from_keypaths
kd.kontext.set_by_path
kd.losses
kd.losses.AbsoluteValue
kd.losses.Huber
kd.losses.L1
kd.losses.L2
kd.losses.Loss
kd.losses.NegativeCosineSimilarity
kd.losses.SigmoidBinaryCrossEntropy
kd.losses.SoftmaxCrossEntropy
kd.losses.SoftmaxCrossEntropyWithIntLabels
kd.losses.Value
kd.losses.compute_losses
kd.metrics
kd.metrics.Accuracy
kd.metrics.Ari
kd.metrics.AutoState
kd.metrics.AverageState
kd.metrics.BinaryAccuracy
kd.metrics.CollectFirstState
kd.metrics.CollectingState
kd.metrics.EmptyState
kd.metrics.LpipsVgg
kd.metrics.Metric
kd.metrics.NoopMetric
kd.metrics.Norm
kd.metrics.Precision1
kd.metrics.Psnr
kd.metrics.RocAuc
kd.metrics.SingleDimension
kd.metrics.SkipIfMissing
kd.metrics.Ssim
kd.metrics.State
kd.metrics.Std
kd.metrics.TreeMap
kd.metrics.TreeReduce
kd.metrics.concat_field
kd.metrics.static_field
kd.metrics.sum_field
kd.metrics.truncate_field
kd.nn
kd.nn.AddEmbedding
kd.nn.AddLearnedEmbedding
kd.nn.AttentionModule
kd.nn.Dropout
kd.nn.DummyModel
kd.nn.ExternalModule
kd.nn.FlatAutoencoder
kd.nn.FourierEmbedding
kd.nn.Identity
kd.nn.ImageTokenizer
kd.nn.ImprovedMultiHeadDotProductAttention
kd.nn.LearnedEmbedding
kd.nn.MultiHeadDotProductAttention
kd.nn.NormModule
kd.nn.ParallelAttentionBlock
kd.nn.Patchify
kd.nn.PatchifyEmbed
kd.nn.PostNormBlock
kd.nn.PreNormBlock
kd.nn.Rearrange
kd.nn.Reduce
kd.nn.Sequential
kd.nn.TransformerBlock
kd.nn.TransformerMLP
kd.nn.Vit
kd.nn.VitEncoder
kd.nn.WrapperModule
kd.nn.ZeroEmbedding
kd.nn.convert_to_fourier_features
kd.nn.interms_property
kd.nn.set_train_property
kd.nn.train_property
kd.optim
kd.optim.decay_to_init
kd.optim.exclude
kd.optim.named_chain
kd.optim.partial_updates
kd.optim.select
kd.random
kd.random.PRNGKey
kd.sharding
kd.summaries
kd.summaries.Histogram
kd.summaries.HistogramSummary
kd.summaries.PointCloud
kd.summaries.ShowBoxes
kd.summaries.ShowDifferenceImages
kd.summaries.ShowImages
kd.summaries.ShowPointCloud
kd.summaries.ShowSegmentations
kd.summaries.ShowTexts
kd.summaries.deprecated
kd.summaries.deprecated.ImageSummary
kd.summaries.deprecated.Summary
kd.testing
kd.testing.assert_step_specs
kd.train
kd.train.Auxiliaries
kd.train.AuxiliariesOutput
kd.train.AuxiliariesState
kd.train.Context
kd.train.RngStream
kd.train.RngStreams
kd.train.Setup
kd.train.TqdmInfo
kd.train.TrainState
kd.train.TrainStep
kd.train.Trainer
kd.train.forward
kd.train.forward_with_loss
kd.typing
kd.typing.Any
kd.typing.Array
kd.typing.ArraySpec
kd.typing.Axes
kd.typing.AxisName
kd.typing.Bool
kd.typing.Callable
kd.typing.Complex
kd.typing.Complex64
kd.typing.DType
kd.typing.Dim
kd.typing.ElementSpec
kd.typing.Float
kd.typing.Float32
kd.typing.Hashable
kd.typing.Initializer
kd.typing.Int
kd.typing.Integer
kd.typing.Memo
kd.typing.Num
kd.typing.PRNGKey
kd.typing.PRNGKeyLike
kd.typing.PyTree
kd.typing.Scalar
kd.typing.ScalarFloat
kd.typing.ScalarInt
kd.typing.Schedule
kd.typing.Sequence
kd.typing.Shape
kd.typing.TfArray
kd.typing.TfFloat
kd.typing.TfFloat32
kd.typing.TfInt
kd.typing.TfUInt8
kd.typing.TypeCheckError
kd.typing.UInt32
kd.typing.UInt8
kd.typing.Union
kd.typing.XArray
kd.typing.check_type
kd.typing.enable_kd_type_checking
kd.typing.set_shape
kd.typing.typechecked
kd.utils
kd.xm
kd.xm.Experiment
kd.xm.WorkUnit
kd.xm.add_colab_artifacts
kd.xm.add_log_artifacts
kd.xm.add_tags_to_xm
kd.xm.load_config_from_path
kd.xm.xmanager_api
.md
.pdf
kd.typing.Axes
kd.typing.Axes
#
[[Source]]
Axes
=
int
|
tuple
[
int
,
...
]