Skip to content

Treex

State

Initializing search

Treex

Treex
User Guide
User Guide
- User Guide
- Defining Modules
- Initialization
- Module API
  Module API
  - Filter
  - Merge
  - Map
- Optimizer
- State Management
- Training State
- Freezing Modules
API Reference
API Reference
- Apply
- ArrayLike
- BatchNorm
- BatchStat
- Cache
- Compact
- Conv
- ConvTranspose
- Copy
- Dropout
- Embed
- Extensions
- FieldInfo
- FieldMetadata
- Filter
- Filters
- Flatten
- FlattenMode
- FlaxModule
- GroupNorm
- HaikuModule
- Hashable
- Initializer
- Inputs
- KeySeq
- KindMixin
- Lambda
- LayerNorm
- Linear
- Log
- Loss
- LossAndLogs
- LossLog
- MISSING
- MLP
- Map
- Merge
- Metric
- MetricLog
- MetricState
- Missing
- ModelState
- Module
- ModuleMeta
- NOTHING
- Named
- Nothing
- Opaque
- OpaquePredicate
- OptState
- Optimizer
- Parameter
- Repr
- Rng
- Sequential
- State State
  Table of contents
  - treex.types.State
- ToDict
- ToString
- Tree
- TreeMeta
- TreePart
- Treex
- add_field_info
- apply
- compact
- compact_module
- copy
- field
- filter
- flatten_mode
- in_compact
- losses
  losses
- map
- merge
- metrics
  metrics
  - Accuracy
  - AuxLosses
  - AuxMetrics
  - LossAndLogs
  - Losses
  - MAE
  - MSE
  - Mean
  - MeanAbsoluteError
  - MeanSquareError
  - Metric
  - Metrics
  - Reduce
  - Reduction
- next_key
- nn
  nn
  - BatchNorm
  - Conv
  - ConvTranspose
  - Dropout
  - Embed
  - Flatten
  - FlaxModule
  - GroupNorm
  - HaikuModule
  - Lambda
  - LayerNorm
  - Linear
  - MLP
  - Sequential
  - sequence
- node
- preserve_state
- regularizers
  regularizers
  - L1
  - L1L2
  - L2
- rng_key
- sequence
- static
- to_dict
- to_string

Table of contents

treex.types.State

treex.State

Source code in treex/types.py

class State(TreePart):
    pass

Previous Sequential

Made with Material for MkDocs