jaxley.optimize.optimizer.TypeOptimizer# class TypeOptimizer(optimizer, optimizer_args, opt_params)[source]# Bases: object optax wrapper which allows different argument values for different params. Parameters: optimizer (Callable) optimizer_args (Dict[str, Any]) opt_params (List[Dict[str, Array]]) init(opt_params)[source]# Initialize the optimizers. Equivalent to optax.optimizers.init(). Parameters: opt_params (List[Dict[str, Array]]) Return type: List update(gradient, opt_state)[source]# Update the optimizers. Equivalent to optax.optimizers.update(). Parameters: gradient (Array) opt_state (List) Return type: Tuple[List, List]
