class_docstr: A PartitionedVariable object allows mutable access to a variable.
  
  PartitionedVariables are identified by a collection (e.g., "batch_stats") and
  a name (e.g., "moving_mean"). The value property gives access to the
  variable's content and can be assigned to for mutation.  Additionally,
  PartitionedVariables enforce logical sharding constraints on both retrieval
  and assignment.
class_docstr_tok: ["A","PartitionedVariable","object","allows","mutable","access","to","a","variable",".","PartitionedVariables","are","identified","by","a","collection","(","e.g.",",","``","batch_stats","''",")","and","a","name","(","e.g.",",","``","moving_mean","''",")",".","The","value","property","gives","access","to","the","variable","'s","content","and","can","be","assigned","to","for","mutation",".","Additionally",",","PartitionedVariables","enforce","logical","sharding","constraints","on","both","retrieval","and","assignment","."]
docstr: Initializes a partitioned variable.
  
  Args:
    scope: The scope in which the variable is stored.
    collection: The collection of the variable (e.g., "params").
    name: The name of the variable (e.g., "dense").
    axes: logical axes name of variable.
    fallback: Fallback behavior if no matching rule is found.
docstr_tok: ["Initializes","a","partitioned","variable",".","Args",":","scope",":","The","scope","in","which","the","variable","is","stored",".","collection",":","The","collection","of","the","variable","(","e.g.",",","``","params","''",")",".","name",":","The","name","of","the","variable","(","e.g.",",","``","dense","''",")",".","axes",":","logical","axes","name","of","variable",".","fallback",":","Fallback","behavior","if","no","matching","rule","is","found","."]
code: def __init__(
      self,
      scope,
      collection: str,
      name: str,
      axes: Optional[Tuple[str, ...]] = None,
      fallback: RulesFallback = RulesFallback.AXIS_IS_UNSHARDED,
  ):
      """Initializes a partitioned variable.
  
      Args:
        scope: The scope in which the variable is stored.
        collection: The collection of the variable (e.g., "params").
        name: The name of the variable (e.g., "dense").
        axes: logical axes name of variable.
        fallback: Fallback behavior if no matching rule is found.
      """
      self.scope = scope
      self.collection = collection
      self.name = name
      self.axes = axes
      self.fallback = fallback
  
imports: import functools
  import re
  from typing import Any, Callable, Mapping, Optional, Tuple
  import flax
  from flax import linen
  from flax import struct
  from flax.core.frozen_dict import freeze
  from flax.core.frozen_dict import unfreeze
  from flax.core.lift import In
  from flax.core.lift import Out
  from flax.linen.spmd import _axis_rules
  from flax.linen.spmd import _AxisRules
  from flax.linen.spmd import _is_logical_spec
  from flax.linen.spmd import _with_sharding_constraint
  from flax.linen.spmd import Array
  from flax.linen.spmd import ArrayPytree
  from flax.linen.spmd import get_logical_axis_rules
  from flax.linen.spmd import logical_axis_rules
  from flax.linen.spmd import logical_to_mesh
  from flax.linen.spmd import logical_to_mesh_axes
  from flax.linen.spmd import LogicalPartitionSpec
  from flax.linen.spmd import LogicalPartitionSpecPytree
  from flax.linen.spmd import LogicalRules
  from flax.linen.spmd import PartitionSpecPytree
  from flax.linen.spmd import RulesFallback
  from flax.linen.spmd import set_logical_axis_rules
  from flax.linen.spmd import with_logical_constraint
  from flax.traverse_util import flatten_dict
  from flax.traverse_util import unflatten_dict
  import jax
label_desc: Use image node_id 1 to create a new PartitionedVariable object from inherited base classes: flax.core.scope with example: obj = PartitionedVariable(scope, collection, name, axes, fallback)