Merge pull request #3792 from chiamp:autosummary

PiperOrigin-RevId: 623230845
google · Apr 9, 2024 · 5ce3485 · 5ce3485
2 parents 2718455 + 1f6cd8a
commit 5ce3485
Show file tree

Hide file tree

Showing 10 changed files with 0 additions and 185 deletions.
diff --git a/docs/api_reference/flax.linen/activation_functions.rst b/docs/api_reference/flax.linen/activation_functions.rst
@@ -33,34 +33,3 @@ Activation functions
 .. autofunction:: standardize
 .. autofunction:: swish
 .. autofunction:: tanh
-
-**Summary**
-
-.. autosummary::
-  :toctree: _autosummary
-
-  PReLU
-  celu
-  elu
-  gelu
-  glu
-  hard_sigmoid
-  hard_silu
-  hard_swish
-  hard_tanh
-  leaky_relu
-  log_sigmoid
-  log_softmax
-  logsumexp
-  one_hot
-  relu
-  relu6 as relu6,
-  selu
-  sigmoid
-  silu
-  soft_sign
-  softmax
-  softplus
-  standardize
-  swish
-  tanh
diff --git a/docs/api_reference/flax.linen/decorators.rst b/docs/api_reference/flax.linen/decorators.rst
@@ -5,11 +5,3 @@ Decorators
 
 .. autofunction:: compact
 .. autofunction:: nowrap
-
-**Summary**
-
-.. autosummary::
-  :toctree: _autosummary
-
-  compact
-  nowrap
diff --git a/docs/api_reference/flax.linen/init_apply.rst b/docs/api_reference/flax.linen/init_apply.rst
@@ -7,13 +7,3 @@ Init/Apply
 .. autofunction:: apply
 .. autofunction:: init
 .. autofunction:: init_with_output
-
-
-**Summary**
-
-.. autosummary::
-  :toctree: _autosummary
-
-  apply
-  init
-  init_with_output
diff --git a/docs/api_reference/flax.linen/initializers.rst b/docs/api_reference/flax.linen/initializers.rst
@@ -25,31 +25,3 @@ Initializers
 .. autofunction:: xavier_uniform
 .. autofunction:: zeros
 .. autofunction:: zeros_init
-
-**Summary**
-
-.. autosummary::
-  :toctree: _autosummary
-
-  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
-  standardize
-  variance_scaling
-  xavier_normal
-  xavier_uniform
-  zeros
-  zeros_init
diff --git a/docs/api_reference/flax.linen/inspection.rst b/docs/api_reference/flax.linen/inspection.rst
@@ -5,10 +5,3 @@ Inspection
 .. currentmodule:: flax.linen
 
 .. autofunction:: tabulate
-
-**Summary**
-
-.. autosummary::
-  :toctree: _autosummary
-
-  tabulate
diff --git a/docs/api_reference/flax.linen/layers.rst b/docs/api_reference/flax.linen/layers.rst
@@ -152,49 +152,3 @@ BatchApply
 .. flax_module::
   :module: flax.linen
   :class: BatchApply
-
-
-**Summary**
-
-.. autosummary::
-  :toctree: _autosummary
-  :template: flax_module
-
-  Dense
-  DenseGeneral
-  Conv
-  ConvTranspose
-  ConvLocal
-  Einsum
-  Embed
-  BatchNorm
-  LayerNorm
-  GroupNorm
-  RMSNorm
-  SpectralNorm
-  WeightNorm
-  Sequential
-  Dropout
-  MultiHeadDotProductAttention
-  MultiHeadAttention
-  SelfAttention
-  RNNCellBase
-  LSTMCell
-  OptimizedLSTMCell
-  SimpleCell
-  GRUCell
-  MGUCell
-  RNN
-  Bidirectional
-  BatchApply
-
-.. autosummary::
-  :toctree: _autosummary
-
-  max_pool
-  avg_pool
-  pool
-  dot_product_attention_weights
-  dot_product_attention
-  make_attention_mask
-  make_causal_mask
diff --git a/docs/api_reference/flax.linen/profiling.rst b/docs/api_reference/flax.linen/profiling.rst
@@ -6,12 +6,3 @@ Profiling
 .. autofunction:: enable_named_call
 .. autofunction:: disable_named_call
 .. autofunction:: override_named_call
-
-**Summary**
-
-.. autosummary::
-  :toctree: _autosummary
-
-  enable_named_call
-  disable_named_call
-  override_named_call
diff --git a/docs/api_reference/flax.linen/spmd.rst b/docs/api_reference/flax.linen/spmd.rst
@@ -18,23 +18,3 @@ SPMD
 .. autofunction:: logical_to_mesh_sharding
 .. autofunction:: with_logical_constraint
 .. autofunction:: with_logical_partitioning
-
-**Summary**
-
-.. autosummary::
-  :toctree: _autosummary
-
-  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
-
diff --git a/docs/api_reference/flax.linen/transformations.rst b/docs/api_reference/flax.linen/transformations.rst
@@ -16,21 +16,3 @@ Transformations
 .. autofunction:: while_loop
 .. autofunction:: cond
 .. autofunction:: switch
-
-.. autosummary::
-  :toctree: _autosummary
-
-  vmap
-  scan
-  jit
-  remat
-  remat_scan
-  map_variables
-  jvp
-  vjp
-  grad
-  value_and_grad
-  custom_vjp
-  while_loop
-  cond
-  switch
diff --git a/docs/api_reference/flax.linen/variable.rst b/docs/api_reference/flax.linen/variable.rst
@@ -4,11 +4,3 @@ Variable dictionary
 
 .. automodule:: flax.core.variables
 .. autoclass:: flax.linen.Variable
-
-**Summary**
-
-.. currentmodule:: flax.linen
-.. autosummary::
-  :toctree: _autosummary
-
-  Variable