import flax.linen as nn import einx import flax from functools import partial import jax.numpy as jnp from typing import Callable, Union, Optional, Any tnn = einx.tracer.import_("flax.linen", "nn") class ParamFactory: class Concrete(einx.tracer.input.Input): def __init__(self, module, name, init, dtype, col, param_type): self.module = module self.name = name self.init = init if dtype is None: if hasattr(module, "dtype"): dtype = module.dtype else: dtype = "float32" self.dtype = dtype self.col = col if param_type == "param": if col is not None: raise ValueError("col is not accepted for flax.linen.Module.param") elif param_type == "variable": if col is None: raise ValueError("col must be specified for flax.linen.Module.variable") else: raise ValueError(f"Unknown tensor factory flax.linen.Module.{param_type}") self.param_type = param_type def to_value_and_key(self): return self.module, ParamFactory.CacheKey( self.name, self.init, self.dtype, self.col, self.param_type ) class CacheKey(einx.tracer.input.CacheKey): def __init__(self, name, init, dtype, col, param_type): self.name = name self.init = init self.dtype = dtype self.col = col self.param_type = param_type def __hash__(self): return hash((self.name, self.init, self.dtype, self.col, self.param_type)) def __eq__(self, other): return ( isinstance(other, ParamFactory.CacheKey) and self.name == other.name and self.init == other.init and self.dtype == other.dtype and self.col == other.col and self.param_type == other.param_type ) def to_tracer(self, backend, virtual_arg): x = ParamFactory.Tracer(self.name, self.init, self.dtype, self.col, self.param_type) return x, x class Tracer(einx.tracer.TensorFactory): def __init__(self, name, init, dtype, col, param_type): self.name = name self.init = init self.dtype = dtype self.col = col self.param_type = param_type def __call__(self, shape, kwargs): name = self.name if not self.name is None else kwargs.get("name", None) init = self.init if not self.init is None else kwargs.get("init", None) dtype = self.dtype if not self.dtype is None else kwargs.get("dtype", None) col = self.col if name is None: raise ValueError( "Must specify name for tensor factory flax.linen.Module.{param|variable}" ) if init is None: raise ValueError( "Must specify init for tensor factory flax.linen.Module.{param|variable}" ) elif isinstance(init, str): if init == "get_at" or init == "rearrange": init = tnn.initializers.normal(stddev=0.02) elif init == "add": init = tnn.initializers.zeros_init() elif init == "multiply": init = tnn.initializers.ones_init() elif init == "dot": init = tnn.initializers.lecun_normal( kwargs["in_axis"], kwargs["out_axis"], kwargs["batch_axis"] ) else: raise ValueError(f"Don't know which initializer to use for operation '{init}'") elif isinstance(init, (int, float)): init = tnn.initializers.constant(init, dtype=dtype) if self.param_type == "param": x = einx.tracer.apply( self.param, args=[name, init, shape, dtype], output=einx.tracer.Tensor(shape) ) else: assert self.param_type == "variable" # Assume that variable initialization does not need an rng key by passing None x = einx.tracer.apply( self.variable, args=[col, name, init, None, shape, dtype], ) x = einx.tracer.apply( einx.tracer.MemberAccess(), args=[x, "value"], output=einx.tracer.Tensor(shape) ) return x def param( x: Union[Callable, nn.Module], name: Optional[str] = None, init: Optional[Any] = None, dtype: Optional[nn.dtypes.Dtype] = None, col: Optional[str] = None, ): """Create a tensor factory for Flax parameters. Args: x: The bound method of a Flax module, i.e. ``nn.Module.param`` or ``nn.Module.variable``, or a module instance in which case its ``param`` method is used. name: Name of the parameter. If ``None``, uses a default name determined from the calling operation. Defaults to ``None``. init: Initializer for the parameter. If ``None``, uses a default init method determined from the calling operation. Defaults to ``None``. dtype: Data type of the parameter. If ``None``, uses the ``dtype`` member of the calling module or ``float32`` if it does not exist. Defaults to ``None``. col: The collection name to use when ``bound_method`` is ``nn.Module.variable``. Returns: A tensor factory with the given default parameters. """ if hasattr(x, "__func__") and x.__func__ == nn.Module.param: module = x.__self__ param_type = "param" elif hasattr(x, "__func__") and x.__func__ == nn.Module.variable: module = x.__self__ param_type = "variable" elif isinstance(x, nn.Module): module = x param_type = "param" else: raise ValueError("x must be a bound method of a Flax module or a Flax module instance") return ParamFactory.Concrete(module, name, init, dtype, col, param_type) # Allow passing nn.Module, nn.Module.param, nn.Module.variable as tensor factory: @einx.tracer.input.register_tensor_factory def tensor_factory(x): if isinstance(x, nn.Module) or ( hasattr(x, "__func__") and (x.__func__ == nn.Module.param or x.__func__ == nn.Module.variable) ): return param(x).to_value_and_key() else: return None # Using _ prefix on classes and a separater constructor, since dataclass/nn.Module does # not support **kwargs parameter. class _Norm(nn.Module): stats: str params: str = "b... [c]" mean: bool = True var: bool = True scale: bool = True bias: bool = True decay_rate: float = None epsilon: float = 1e-5 fastvar: bool = True dtype: nn.dtypes.Dtype = "float32" kwargs: dict = None @nn.compact def __call__(self, x, training=None): if self.decay_rate is not None and training is None: raise ValueError("training must be specified when decay_rate is used") use_ema = self.decay_rate is not None and (not training or self.is_initializing()) x, mean, var = einx.nn.norm( x, self.stats, self.params, mean=param(self.variable, col="stats", name="mean", dtype=self.dtype) if use_ema and self.mean else self.mean, var=param(self.variable, col="stats", name="var", dtype=self.dtype) if use_ema and self.var else self.var, scale=param(self.param, name="scale", dtype=self.dtype) if self.scale else None, bias=param(self.param, name="bias", dtype=self.dtype) if self.bias else None, epsilon=self.epsilon, fastvar=self.fastvar, **(self.kwargs if self.kwargs is not None else {}), ) update_ema = self.decay_rate is not None and training and not self.is_initializing() if update_ema: if self.mean: mean_ema = self.variable("stats", "mean", None) mean_ema.value = self.decay_rate * mean_ema.value + (1 - self.decay_rate) * mean if self.var: var_ema = self.variable("stats", "var", None) var_ema.value = self.decay_rate * var_ema.value + (1 - self.decay_rate) * var return x def Norm( stats: str, params: str = "b... [c]", mean: bool = True, var: bool = True, scale: bool = True, bias: bool = True, decay_rate: Optional[float] = None, epsilon: float = 1e-5, fastvar: bool = True, dtype: nn.dtypes.Dtype = "float32", name: Optional[str] = None, **kwargs: Any, ): """Normalization layer. Args: stats: Einstein string determining the axes along which mean and variance are computed. Will be passed to ``einx.reduce``. params: Einstein string determining the axes along which learnable parameters are applied. Will be passed to ``einx.elementwise``. Defaults to ``"b... [c]"``. mean: Whether to apply mean normalization. Defaults to ``True``. var: Whether to apply variance normalization. Defaults to ``True``. scale: Whether to apply a learnable scale according to ``params``. Defaults to ``True``. bias: Whether to apply a learnable bias according to ``params``. Defaults to ``True``. epsilon: A small float added to the variance to avoid division by zero. Defaults to ``1e-5``. fastvar: Whether to use a fast variance computation. Defaults to ``True``. dtype: Data type of the weights. Defaults to ``"float32"``. decay_rate: Decay rate for exponential moving average of mean and variance. If ``None``, no moving average is applied. Defaults to ``None``. name: Name of the module. Defaults to ``None``. **kwargs: Additional parameters that specify values for single axes, e.g. ``a=4``. """ return _Norm( stats, params=params, mean=mean, var=var, scale=scale, bias=bias, decay_rate=decay_rate, epsilon=epsilon, fastvar=fastvar, dtype=dtype, name=name, kwargs=kwargs, ) class _Linear(nn.Module): expr: str bias: bool = True dtype: nn.dtypes.Dtype = "float32" kwargs: dict = None @nn.compact def __call__(self, x): return einx.nn.linear( x, self.expr, bias=param(self.param, name="bias", dtype=self.dtype) if self.bias else None, weight=param(self.param, name="weight", dtype=self.dtype), **(self.kwargs if self.kwargs is not None else {}), ) def Linear( expr: str, bias: bool = True, dtype: nn.dtypes.Dtype = "float32", name: Optional[str] = None, **kwargs: Any, ): """Linear layer. Args: expr: Einstein string determining the axes along which the weight matrix is multiplied. Will be passed to ``einx.dot``. bias: Whether to apply a learnable bias. Defaults to ``True``. dtype: Data type of the weights. Defaults to ``"float32"``. name: Name of the module. Defaults to ``None``. **kwargs: Additional parameters that specify values for single axes, e.g. ``a=4``. """ return _Linear(expr, bias=bias, dtype=dtype, name=name, kwargs=kwargs) class _Dropout(nn.Module): expr: str drop_rate: float rng_collection: str = "dropout" kwargs: dict = None @nn.compact def __call__(self, x, training): if training: return einx.nn.dropout( x, self.expr, drop_rate=self.drop_rate, rng=self.make_rng(self.rng_collection), **(self.kwargs if self.kwargs is not None else {}), ) else: return x def Dropout( expr: str, drop_rate: float, rng_collection: str = "dropout", name: Optional[str] = None, **kwargs: Any, ): """Dropout layer. Args: expr: Einstein string determining the axes along which dropout is applied. Will be passed to ``einx.elementwise``. drop_rate: Drop rate. rng_collection: the rng collection name to use when requesting an rng key. Defaults to ``"dropout"``. name: Name of the module. Defaults to ``None``. **kwargs: Additional parameters that specify values for single axes, e.g. ``a=4``. """ return _Dropout(expr, drop_rate, rng_collection=rng_collection, name=name, kwargs=kwargs)