API 参考

• flax.config 包
  • 配置
    • Config.flax_filter_frames
    • Config.flax_fix_rng_separator
    • Config.flax_preserve_adopted_names
    • Config.flax_profile
    • Config.flax_return_frozendict
    • Config.flax_use_orbax_checkpointing
    • Config.update()
• flax.core.frozen_dict 包
  • FrozenDict
    • FrozenDict.copy()
    • FrozenDict.pop()
    • FrozenDict.pretty_repr()
    • FrozenDict.unfreeze()
  • freeze()
  • unfreeze()
  • copy()
  • pop()
  • pretty_repr()
• flax.cursor 包
  • cursor()
  • Cursor
    • Cursor.apply_update()
    • Cursor.build()
    • Cursor.find()
    • Cursor.find_all()
    • Cursor.set()
• flax.errors 包
  • AlreadyExistsError
  • ApplyModuleInvalidMethodError
  • ApplyScopeInvalidVariablesStructureError
  • ApplyScopeInvalidVariablesTypeError
  • AssignSubModuleError
  • CallCompactUnboundModuleError
  • CallSetupUnboundModuleError
  • CallShareScopeOnUnboundModuleError
  • CallUnbindOnUnboundModuleError
  • CursorFindError
  • DescriptorAttributeError
  • IncorrectPostInitOverrideError
  • InvalidCheckpointError
  • InvalidFilterError
  • InvalidInstanceModuleError
  • InvalidRngError
  • InvalidScopeError
  • JaxTransformError
  • LazyInitError
  • MPACheckpointingRequiredError
  • MPARestoreDataCorruptedError
  • MPARestoreTargetRequiredError
  • ModifyScopeVariableError
  • MultipleMethodsCompactError
  • NameInUseError
  • PartitioningUnspecifiedError
  • ReservedModuleAttributeError
  • ScopeCollectionNotFound
  • ScopeParamNotFoundError
  • ScopeParamShapeError
  • ScopeVariableNotFoundError
  • SetAttributeFrozenModuleError
  • SetAttributeInModuleSetupError
  • TraceContextError
  • TransformTargetError
  • TransformedMethodReturnValueError
  • TraverseTreeError
• flax.jax_utils 包
  • partial_eval_by_shape()
  • 多设备实用程序
    • replicate()
    • unreplicate()
    • prefetch_to_device()
    • pmean()
    • pad_shard_unpad()
• flax.linen
  • 模块
    • 模块
      • Module.__setattr__()
      • Module.apply()
      • Module.bind()
      • Module.copy()
      • Module.get_variable()
      • Module.has_rng()
      • Module.has_variable()
      • Module.init()
      • Module.init_with_output()
      • Module.is_initializing()
      • Module.is_mutable_collection()
      • Module.lazy_init()
      • Module.make_rng()
      • Module.module_paths()
      • Module.param()
      • Module.path
      • Module.perturb()
      • Module.put_variable()
      • Module.setup()
      • Module.sow()
      • Module.tabulate()
      • Module.unbind()
      • Module.variable()
      • Module.variables
    • apply()
    • init()
    • init_with_output()
    • intercept_methods()
    • share_scope()
  • Init/Apply
    • apply()
    • init()
    • init_with_output()
  • 层
    • 线性模块
      • Dense
      • DenseGeneral
      • Conv
      • ConvTranspose
      • ConvLocal
      • Einsum
      • Embed
    • 池化
      • max_pool()
      • avg_pool()
      • pool()
    • 归一化
      • BatchNorm
      • LayerNorm
      • GroupNorm
      • RMSNorm
      • InstanceNorm
      • SpectralNorm
      • WeightNorm
    • 组合器
      • Sequential
    • 随机
      • Dropout
    • 注意力
      • MultiHeadDotProductAttention
      • MultiHeadAttention
      • SelfAttention
      • dot_product_attention_weights()
      • dot_product_attention()
      • make_attention_mask()
      • make_causal_mask()
    • 循环
      • RNNCellBase
      • LSTMCell
      • OptimizedLSTMCell
      • ConvLSTMCell
      • SimpleCell
      • GRUCell
      • MGUCell
      • RNN
      • Bidirectional
    • BatchApply
      • BatchApply
  • 激活函数
    • PReLU
      • PReLU.param_dtype
      • PReLU.negative_slope_init
      • PReLU.__call__()
      • PReLU.param_dtype
    • celu()
    • elu()
    • gelu()
    • glu()
    • hard_sigmoid()
    • hard_silu()
    • hard_swish()
    • hard_tanh()
    • leaky_relu()
    • log_sigmoid()
    • log_softmax()
    • logsumexp()
    • one_hot()
    • relu()
    • selu()
    • sigmoid()
    • silu()
    • soft_sign()
    • softmax()
    • softplus()
    • standardize()
    • swish()
    • tanh()
  • 初始化器
    • constant()
    • delta_orthogonal()
    • glorot_normal()
    • glorot_uniform()
    • he_normal()
    • he_uniform()
    • kaiming_normal()
    • kaiming_uniform()
    • lecun_normal()
    • lecun_uniform()
    • normal()
    • truncated_normal()
    • ones()
    • ones_init()
    • orthogonal()
    • uniform()
    • variance_scaling()
    • xavier_normal()
    • xavier_uniform()
    • zeros()
    • zeros_init()
  • 转换
    • vmap()
    • scan()
    • jit()
    • remat()
    • remat_scan()
    • map_variables()
    • jvp()
    • vjp()
    • custom_vjp()
    • while_loop()
    • cond()
    • switch()
  • 检查
    • tabulate()
  • 变量字典
    • 变量
  • SPMD
    • Partitioned()
    • with_partitioning()
    • get_partition_spec()
    • get_sharding()
    • LogicallyPartitioned()
    • logical_axis_rules()
    • set_logical_axis_rules()
    • get_logical_axis_rules()
    • logical_to_mesh_axes()
    • logical_to_mesh()
    • logical_to_mesh_sharding()
    • with_logical_constraint()
    • with_logical_partitioning()
  • 装饰器
    • compact()
    • nowrap()
  • 性能分析
    • enable_named_call()
    • disable_named_call()
    • override_named_call()
• flax.serialization 包
  • 状态字典
    • from_state_dict()
    • to_state_dict()
    • register_serialization_state()
  • 使用 MessagePack 进行序列化
    • msgpack_serialize()
    • msgpack_restore()
    • to_bytes()
    • from_bytes()
• flax.struct 包
  • dataclass()
  • PyTreeNode
• flax.traceback_util 包
  • 回溯过滤实用程序
    • hide_flax_in_tracebacks()
    • show_flax_in_tracebacks()
• flax.training 包
  • 检查点
    • save_checkpoint()
    • save_checkpoint_multiprocess()
    • latest_checkpoint()
    • restore_checkpoint()
    • convert_pre_linen()
  • 学习率计划
    • create_constant_learning_rate_schedule()
    • create_stepped_learning_rate_schedule()
    • create_cosine_learning_rate_schedule()
  • 训练状态
    • TrainState
      • TrainState.apply_gradients()
      • TrainState.create()
  • 提前停止
    • EarlyStopping
      • EarlyStopping.min_delta
      • EarlyStopping.patience
      • EarlyStopping.best_metric
      • EarlyStopping.patience_count
      • EarlyStopping.should_stop
      • EarlyStopping.has_improved
      • EarlyStopping.update()
  • 常用实用程序
    • shard()
    • shard_prng_key()
    • stack_forest()
    • get_metrics()
    • onehot()
• flax.traverse_util 包
  • 遍历对象
    • Traversal
      • Traversal.compose()
      • Traversal.each()
      • Traversal.filter()
      • Traversal.iterate()
      • Traversal.merge()
      • Traversal.set()
      • Traversal.tree()
      • Traversal.update()
    • TraverseId
      • TraverseId.iterate()
      • TraverseId.update()
    • TraverseMerge
      • TraverseMerge.iterate()
      • TraverseMerge.update()
    • TraverseCompose
      • TraverseCompose.iterate()
      • TraverseCompose.update()
    • TraverseFilter
      • TraverseFilter.iterate()
      • TraverseFilter.update()
    • TraverseAttr
      • TraverseAttr.iterate()
      • TraverseAttr.update()
    • TraverseItem
      • TraverseItem.iterate()
      • TraverseItem.update()
    • TraverseEach
      • TraverseEach.iterate()
      • TraverseEach.update()
    • TraverseTree
      • TraverseTree.iterate()
      • TraverseTree.update()
  • 字典实用程序
    • flatten_dict()
    • unflatten_dict()
    • path_aware_map()
  • 模型参数遍历
    • ModelParamTraversal
      • ModelParamTraversal.__init__()