def _get_loss(self, loss_name: str): if loss_name == 'logistic': x = TrainVar(jn.zeros(2)) model_vars = VarCollection({'x': x}) def loss(): return jn.log(jn.exp(-jn.sum(x.value)) + 1) return model_vars, loss if loss_name == 'square': # loss = x*x + y*y. x = TrainVar(jn.ones(2)) y = TrainVar(jn.ones(3)) model_vars = VarCollection({'x': x, 'y': y}) def loss(): return jn.dot(x.value, x.value) + jn.dot(y.value, y.value) return model_vars, loss if loss_name == 'rastrigin': d = 2 x = TrainVar(jn.ones(d)) model_vars = VarCollection({'x': x}) def loss(): return 10 * d + jn.dot(x.value, x.value) - 10 * jn.sum( jn.cos(2 * math.pi * x.value)) return model_vars, loss raise ValueError
def __init__(self, f: Callable, vc: VarCollection): """Function constructor. Args: f: the function or the module to represent. vc: the VarCollection of variables used by the function. """ if hasattr(f, '__name__'): self.vc = VarCollection((f'{{{f.__name__}}}{k}', v) for k, v in vc.items()) else: self.vc = VarCollection(vc) self.__wrapped__ = f
def __init__(self, f: Callable, variables: Optional[VarCollection], input_argnums: Optional[Tuple[int, ...]] = None): """Constructs an instance to compute the gradient of f w.r.t. variables. Args: f: the function for which to compute gradients. variables: the variables for which to compute gradients. input_argnums: input indexes, if any, on which to compute gradients. """ variables = variables or VarCollection() super().__init__(variables) self.input_argnums = input_argnums or tuple() def f_func(inputs_and_train_tensors: List[JaxArray], state_tensors: List[JaxArray], list_args: List): inputs = inputs_and_train_tensors[:len(self.input_argnums)] train_tensors = inputs_and_train_tensors[len(self.input_argnums):] original_vc = self.vc.tensors() self.vc.subset(TrainVar).assign(train_tensors) self.vc.subset(BaseState).assign(state_tensors) for i, arg in zip(self.input_argnums, inputs): list_args[i] = arg outputs = f(*list_args) if not isinstance(outputs, (list, tuple)): outputs = [outputs] return_value = outputs[0], (outputs, variables.tensors(BaseState)) self.vc.assign(original_vc) return return_value self.f = jax.grad(f_func, has_aux=True)
def vars(self, scope: str = '') -> VarCollection: """Collect all the variables (and their names) contained in the VarCollection. Args: scope: string to prefix to the variable names. Returns: A VarCollection of all the variables. """ return VarCollection((scope + k, v) for k, v in self.vc.items())
def vars(self, scope: str = '') -> VarCollection: """Collect all the variables (and their names) contained in the list and its submodules. Args: scope: string to prefix to the variable names. Returns: A VarCollection of all the variables. """ vc = VarCollection() scope += f'({self.__class__.__name__})' for p, v in enumerate(self): if isinstance(v, BaseVar): vc[f'{scope}[{p}]'] = v elif isinstance(v, Module): vc.update(v.vars(scope=f'{scope}[{p}]')) return vc
def __init__(self, derivative_fn: Callable, f: Union[Module, Callable], variables: Optional[VarCollection], input_argnums: Optional[Tuple[int, ...]] = None, return_all_f_outputs: bool = False): """Constructs an instance to compute the derivatives of f w.r.t. variables. Args: derivative_fn: JAX transformation which computes derivative. f: the function for which to compute derivatives. variables: the variables for which to compute derivatives. input_argnums: input indexes, if any, on which to compute derivatives. return_all_f_outputs: if True also return original outputs of the fuction along with derivatives. """ self.f = f self.vc = variables = VarCollection(variables or ()) if not isinstance(f, Module): f = Function(f, self.vc) assert isinstance(input_argnums, tuple) or input_argnums is None, \ f'Must pass a tuple of indices to input_argnums; received {input_argnums}.' self.input_argnums = input_argnums or tuple() self.return_all_f_outputs = return_all_f_outputs def f_func(inputs_and_train_tensors: List[JaxArray], state_tensors: List[JaxArray], list_args: List, kwargs: Dict): inputs = inputs_and_train_tensors[:len(self.input_argnums)] train_tensors = inputs_and_train_tensors[len(self.input_argnums):] original_vc = self.vc.tensors() try: self.vc.subset(TrainVar).assign(train_tensors) self.vc.subset(BaseState).assign(state_tensors) for i, arg in zip(self.input_argnums, inputs): list_args[i] = arg outputs = f(*list_args, **kwargs) if not isinstance(outputs, (list, tuple)): outputs = [outputs] if self.return_all_f_outputs: return outputs[0], (outputs, variables.tensors()) else: return outputs[0], variables.tensors() finally: self.vc.assign(original_vc) self._call = derivative_fn(f_func)
def vars(self, scope: str = '') -> VarCollection: """Collect all the variables (and their names) contained in the module and its submodules. Important: Variables and modules stored Python structures such as dict or list are not collected. See ModuleList if you need such a feature. Args: scope: string to prefix to the variable names. Returns: A VarCollection of all the variables. """ vc = VarCollection() scope += f'({self.__class__.__name__}).' for k, v in self.__dict__.items(): if isinstance(v, BaseVar): vc[scope + k] = v elif isinstance(v, Module): vc.update(v.vars(scope=scope + k)) return vc
def __init__(self, f: Union[Module, Callable], variables: Optional[VarCollection], input_argnums: Optional[Tuple[int, ...]] = None): """Constructs an instance to compute the gradient of f w.r.t. variables. Args: f: the function for which to compute gradients. variables: the variables for which to compute gradients. input_argnums: input indexes, if any, on which to compute gradients. """ self.f = f self.vc = variables = VarCollection(variables or ()) if not isinstance(f, Module): f = Function(f, self.vc) def f_func(inputs_and_train_tensors: List[JaxArray], state_tensors: List[JaxArray], list_args: List, kwargs: Dict): inputs = inputs_and_train_tensors[:len(self.input_argnums)] train_tensors = inputs_and_train_tensors[len(self.input_argnums):] original_vc = self.vc.tensors() try: self.vc.subset(TrainVar).assign(train_tensors) self.vc.subset(BaseState).assign(state_tensors) for i, arg in zip(self.input_argnums, inputs): list_args[i] = arg outputs = f(*list_args, **kwargs) if not isinstance(outputs, (list, tuple)): outputs = [outputs] return outputs[0], (outputs, variables.tensors()) finally: self.vc.assign(original_vc) assert isinstance(input_argnums, tuple) or input_argnums is None, \ f'Must pass a tuple of indices to input_argnums; received {input_argnums}.' self.input_argnums = input_argnums or tuple() signature = inspect.signature(f) self.__wrapped__ = f self.__signature__ = signature.replace( return_annotation=Tuple[List[JaxArray], signature.return_annotation]) self._call = jax.grad(f_func, has_aux=True)
def __init__(self, closure_vars, global_vars): self.closure_vars = closure_vars self.global_vars = global_vars self.vc = VarCollection()
def vars(self, scope: str = '') -> VarCollection: """Return the VarCollection of the variables used.""" if scope: return VarCollection((scope + k, v) for k, v in self.vc.items()) return VarCollection(self.vc)
def __init__(self, vc: VarCollection): super().__init__() self.vc = VarCollection( (f'({self.__class__.__name__}){k}', v) for k, v in vc.items())