API 参考

• flax.nnx
  • graph
    • split()
    • merge()
    • update()
    • pop()
    • state()
    • graphdef()
    • iter_graph()
    • clone()
    • call()
    • GraphDef
    • UpdateContext
      • UpdateContext.merge()
      • UpdateContext.split()
    • update_context()
    • current_update_context()
  • module
    • Module
      • Module.sow()
      • Module.iter_modules()
      • Module.eval()
      • Module.iter_children()
      • Module.iter_modules()
      • Module.set_attributes()
      • Module.sow()
      • Module.train()
  • nn
    • 激活函数
      • 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()
    • 注意力
      • MultiHeadAttention
      • combine_masks()
      • dot_product_attention()
      • make_attention_mask()
      • make_causal_mask()
    • 初始化器
      • 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()
    • 线性
      • Conv
      • ConvTranspose
      • Embed
      • 线性
      • LinearGeneral
      • Einsum
    • 规范化
      • BatchNorm
      • LayerNorm
      • RMSNorm
      • GroupNorm
    • 随机
      • Dropout
  • rnglib
    • Rngs
      • Rngs.__init__()
    • RngStream
    • reseed()
  • spmd
    • get_partition_spec()
    • get_named_sharding()
    • with_partitioning()
    • with_sharding_constraint()
  • state
    • State
      • State.filter()
      • State.merge()
      • State.split()
  • training
    • 指标
      • Metric
      • Average
      • Accuracy
      • Welford
      • MultiMetric
    • 优化器
      • 优化器
  • transforms
    • Jit
    • Remat
    • Scan
    • Vmap
    • grad()
    • jit()
    • remat()
    • scan()
    • value_and_grad()
    • vmap()
    • eval_shape()
    • cond()
  • variables
    • BatchStat
    • Cache
    • Intermediate
    • Param
    • Variable
    • VariableMetadata
    • VariableState
    • with_metadata()
  • helpers
    • Dict
    • Sequential
    • TrainState
      • TrainState.replace()
  • visualization
    • display()
  • filterlib
    • to_predicate()
    • WithTag
    • PathContains
    • OfType
    • Any
    • All
    • Not
    • Everything
    • Nothing
  • bridge
    • ToNNX
      • ToNNX.__call__()
      • ToNNX.lazy_init()
    • ToLinen
      • ToLinen.__call__()
    • to_linen()
    • NNXMeta
      • NNXMeta.__call__()
      • NNXMeta.add_axis()
      • NNXMeta.get_partition_spec()
      • NNXMeta.remove_axis()
      • NNXMeta.replace()
      • NNXMeta.replace_boxed()
      • NNXMeta.to_nnx_variable()
      • NNXMeta.unbox()
• flax.core.frozen_dict 包
  • FrozenDict
    • FrozenDict.copy()
    • FrozenDict.pop()
    • FrozenDict.pretty_repr()
    • FrozenDict.unfreeze()
  • freeze()
  • unfreeze()
  • copy()
  • pop()
  • pretty_repr()
• flax.struct 包
  • dataclass()
  • PyTreeNode
• flax.training 包
  • 训练状态
    • TrainState
      • TrainState.apply_gradients()
      • TrainState.create()