def assign(self, tensor: JaxArray, check=True): """Sets the value of the variable.""" if check: assert_assigned_type_and_shape_match(self.value, tensor) self.value = tensor
def assign(self, tensor: JaxArray, check=True): if check: assert_assigned_type_and_shape_match(self.value, tensor) self._value = tensor