Prompt:
 optim = getattr(self, f'optim{unet_index}')

        return optim.param_groups[0]['lr']

    # function for allowing only one unet from being trained at a time

    def validate_and_set_unet_being_trained(self, unet_number = None):
        if exists(unet_number):
            self.validate_unet_number(unet_number)

        assert not exists(self.only_train_unet_number) or self.only_train_unet_number == unet_number, 'you cannot only train on one unet at a time. you will need to save the trainer into a checkpoint, and resume training on a new unet'

        self.only_train_unet_number = unet_number
        self.imagen.only_train_unet_number = unet_number

        if not exists(unet_number):
            return

        self.wrap_unet(unet_number)

    def wrap_unet(self, unet_number):
        if hasattr(self, 'one_unet_wrapped'):
            return

        unet = self.imagen.get_unet(unet_number)
        unet_index = unet_number - 1

        optimizer = getattr(self, f'optim{unet_index}')
        scheduler = getattr(self, f'scheduler{unet_index}')

        if self.train_dl:
            self.unet_being_trained, self.train_dl, optimizer = self.accelerator.prepare(unet, self.train_dl, optimizer)
        else:
            self.unet_being_trained, optimizer = self.accelerator.prepare(unet, optimizer)

        if exists(scheduler):
            scheduler = self.accelerator.prepare(scheduler)

        setattr(self, f'optim{unet_index}', optimizer)
        setattr(self, f'scheduler{unet_index}', scheduler)

        self.one_unet_wrapped = True

    # hacking accelerator due to not having separate gradscaler per optimizer

    def set_accelerator_scaler(self, unet_number):
        unet_number = self.validate_unet_number(unet_number)
        scaler = getattr(self, f'scaler{unet_number - 1}')

        self.accelerator.scaler = scaler
        for optimizer in self.accelerator._optimizers:
            optimizer.scaler = scaler

    # helper print

    def print(self, msg):
        if not self.is_main:
            return

        if not self.verbose:
            return

        return self.accelerator.print(msg)

    # validating the unet number

    def validate_unet_number(self, unet_number = None):
        if self.num_unets == 1:
            unet_number = default(unet_number, 1)

        assert 0 < unet_number <= self.num_unets, f'unet number should be in between 1 and {self.num_unets}'
        return unet_number

    # number of training steps taken

    def num_steps_taken(self, unet_number = None):
        if self.num_unets == 1:
            unet_number = default(unet_number, 1)

        return self.steps[unet_number - 1].item()

    def print_untrained_unets(self):
        print_final_error = False

        for ind, (steps, unet) in enumerate(zip(self.steps.tolist(), self.imagen.unets)):
            if steps > 0 or isinstance(unet, NullUnet):
                continue

            self.print(f'unet {ind + 1} has not been trained')
            print_final_error = True

        if print_final_error:
            self.print('when sampling, you can pass stop_at_unet_number to stop early in the cascade, so it does not try to generate with untrained unets')

    # data related functions

    def add_train_dataloader(self, dl = None):
        if not exists(dl):
            return

        assert not exists(self.train_dl), 'training dataloader was already added'
        assert not self.prepared, f'You need to add the dataset before preperation'
        self.train_dl = dl

    def add_valid_dataloader(self, dl):
        if not exists(dl):
            return

        assert not exists(self.valid_dl), 'validation dataloader was already added'
        assert not self.prepared, f'You need to add the dataset before preperation'
        self.valid_dl = dl

    def add_train_dataset(self, ds = None, *, batch_size, **dl_kwargs):
        if not exists(ds):
            return

        assert not exists(self.train_dl), 'training dataloader was already added'

        valid_ds = None
        if self.split_valid_from_train:
            train_size = int((1 - self.split_valid_fraction) * len(ds))
            valid_size = len(ds) - train_size

            ds, valid_ds = random_split(ds, [train_size, valid_size], generator = torch.Generator().manual_seed(self.split_random_seed))
            self.print(f'training with dataset of {len(ds)} samples and validating with randomly splitted {len(valid_ds)} samples')

        dl = DataLoader(ds, batch_size = batch_size, **dl_kwargs)
        self.add_train_dataloader(dl)

        if not self.split_valid_from_train:
            return

        self.add_valid_dataset(valid_ds, batch_size = batch_size, **dl_kwargs)

    def add_valid_dataset(self, ds, *, batch_size, **dl_kwargs):
        if not exists(ds):
            return

        assert not exists(self.valid_dl), 'validation dataloader was already added'

        dl = DataLoader(ds, batch_size = batch_size, **dl_kwargs)
        self.add_valid_dataloader(dl)

    def create_train_iter(self):
        assert exists(self.train_dl), 'training dataloader has not been registered with the trainer yet'

        if exists(self.train_dl_iter):
            return

        self.train_dl_iter = cycle(self.train_dl)

    def create_valid_iter(self):
        assert exists(self.valid_dl), 'validation dataloader has not been registered with the trainer yet'

        if exists(self.valid_dl_iter):
            return

        self.valid_dl_iter = cycle(self.valid_dl)

    def train_step(self, unet_number = None, **kwargs):
        if not self.prepared:
            self.prepare()
        self.create_train_iter()
        loss = self.step_with_dl_iter(self.train_dl_iter, unet_number = unet_number, **kwargs)
        self.update(unet_number = unet_number)
        return loss

    @torch.no_grad()
    @eval_decorator
    def valid_step(self, **kwargs):
        if not self.prepared:
            self.prepare()
        self.create_valid_iter()
        context = self.use_ema_unets if kwargs.pop('use_ema_unets', False) else nullcontext
        with context():
            loss = self.step_with_dl_iter(self.valid_dl_iter, **kwargs)
        return loss

    def step_with_dl_iter(self, dl_iter, **kwargs):
        dl_tuple_output = cast_tuple(next(dl_iter))
        model_input = dict(list(zip(self.dl_tuple_output_keywords_names, dl_tuple_output)))
        loss = self.forward(**{**kwargs, **model_input})
        return loss

    # checkpointing functions

    @property
    def all_checkpoints_sorted(self):
        glob_pattern = os.path.join(self.checkpoint_path, '*.pt')
        checkpoints = self.fs.glob(glob_pattern)
        sorted_checkpoints = sorted(checkpoints, key = lambda x: int(str(x).split('.')[-2]), reverse = True)
        return sorted_checkpoints

    def load_from_checkpoint_folder(self, last_total_steps = -1):
        if last_total_steps != -1:
            filepath = os.path.join(self.checkpoint_path, f'checkpoint.{last_total_steps}.pt')
            self.load(filepath)
            return

        sorted_checkpoints = self.all_checkpoints_sorted

        if len(sorted_checkpoints) == 0:
            self.print(f'no checkpoints found to load from at {self.checkpoint_path}')
            return

        last_checkpoint = sorted_checkpoints[0]
        self.load(last_checkpoint)

    def save_to_checkpoint_folder(self):
        self.accelerator.wait_for_everyone()

        if not self.can_checkpoint:
            return

        total_steps = int(self.steps.sum().item())
        filepath = os.path.join(self.checkpoint_path, f'checkpoint.{total_steps}.pt')

        self.save(filepath)

        if self.max_checkpoints_keep <= 0:
            return

        sorted_checkpoints = self.all_checkpoints_sorted
        checkpoints_to_discard = sorted_checkpoints[self.max_checkpoints_keep:]

        for checkpoint in checkpoints_to_discard:
            self.fs.rm(checkpoint)

    # saving and loading functions

    def save(
        self,
        path,
        overwrite = True,
        without_optim_and_sched = False,
        **kwargs
    ):
        self.accelerator.wait_for_everyone()

        if not self.can_checkpoint:
            return

        fs = self.fs

        assert not (fs.exists(path) and not overwrite)

        self.reset_ema_unets_all_one_device()

        save_obj = dict(
            model = self.imagen.state_dict(),
            version = __version__,
            steps = self.steps.cpu(),
            **kwargs
        )

        save_optim_and_sched_iter = range(0, self.num_unets) if not without_optim_and_sched else tuple()

        for ind in save_optim_and_sched_iter:
            scaler_key = f'scaler{ind}'
            optimizer_key = f'optim{ind}'
            scheduler_key = f'scheduler{ind}'
            warmup_scheduler_key = f'warmup{ind}'

            scaler = getattr(self, scaler_key)
            optimizer = getattr(self, optimizer_key)
            scheduler = getattr(self, scheduler_key)
            warmup_scheduler = getattr(self, warmup_scheduler_key)

            if exists(scheduler):
                save_obj = {**save_obj, scheduler_key: scheduler.state_dict()}

            if exists(warmup_scheduler):
                save_obj = {**save_obj, warmup_scheduler_key: warmup_scheduler.state_dict()}

            save_obj = {**save_obj, scaler_key: scaler.state_dict(), optimizer_key: optimizer.state_dict()}

        if self.use_ema:
            save_obj = {**save_obj, 'ema': self.ema_unets.state_dict()}

        # determine if imagen config is available

        if hasattr(self.imagen, '_config'):
            self.print(f'this checkpoint is commandable from the CLI - "imagen --model {str(path)} \"<prompt>\""')

            save_obj = {
                **save_obj,
                'imagen_type': 'elucidated' if self.is_elucidated else 'original',
                'imagen_params': self.imagen._config
            }

        #save to path

        with fs.open(path, 'wb') as f:
            torch.save(save_obj, f)

        self.print(f'checkpoint saved to {path}')

    def load(self, path, only_model = False, strict = True, noop_if_not_exist = False):
        fs = self.fs

        if noop_if_not_exist and not fs.exists(path):
            self.print(f'trainer checkpoint not found at {str(path)}')
            return

        assert fs.exists(path), f'{path} does not exist'

        self.reset_ema_unets_all_one_device()

        # to avoid extra GPU memory usage in main process when using Accelerate

        with fs.open(path) as f:
            loaded_obj = torch.load(f, map_location='cpu')

        if version.parse(__version__) != version.parse(loaded_obj['version']):
            self.print(f'loading saved imagen at version {loaded_obj["version"]}, but current package version is {__version__}')

        try:
            self.imagen.load_state_dict(loaded_obj['model'], strict = strict)
        except RuntimeError:
            print("Failed loading state dict. Trying partial load")
            self.imagen.load_state_dict(restore_parts(self.imagen.state_dict(),
                                                      loaded_obj['model']))

        if only_model:
            return loaded_obj

        self.steps.copy_(loaded_obj['steps'])

        for ind in range(0, self.num_unets):
            scaler_key = f'scaler{ind}'
            optimizer_key = f'optim{ind}'
            scheduler_key = f'scheduler{ind}'
            warmup_scheduler_key = f'warmup{ind}'

            scaler = getattr(self, scaler_key)
            optimizer = getattr(self, optimizer_key)
            scheduler = getattr(self, scheduler_key)
            warmup_scheduler = getattr(self, warmup_scheduler_key)

            if exists(scheduler) and scheduler_key in loaded_obj:
                scheduler.load_state_dict(loaded_obj[scheduler_key])

            if exists(warmup_scheduler) and warmup_scheduler_key in loaded_obj:
                warmup_scheduler.load_state_dict(loaded_obj[warmup_scheduler_key])

            if exists(optimizer):
                try:
                    optimizer.load_state_dict(loaded_obj[optimizer_key])
                    scaler.load_state_dict(loaded_obj[scaler_key])
                except:
                    self.print('could not load optimizer and scaler, possibly because you have turned on mixed precision training since the last run. resuming with new optimizer and scalers')

        if self.use_ema:
            assert 'ema' in loaded_obj
            try:
                self.ema_unets.load_state_dict(loaded_obj['ema'], strict = strict)
            except RuntimeError:
                print("Failed loading state dict. Trying partial load")
                self.ema_unets.load_state_dict(restore_parts(self.ema_unets.state_dict(),
                                                             loaded_obj['ema']))

        self.print(f'checkpoint loaded from {path}')
        return loaded_obj

    # managing ema unets and their devices

    @property
    def unets(self):
        return nn.ModuleList([ema.ema_model for ema in self.ema_unets])

    def get_ema_unet(self, unet_number = None):

        if not self.use_ema:
            return

        unet_number = self.validate_unet_number(unet_number)
        index = unet_number - 1

        if isinstance(self.unets, nn.ModuleList):
            unets_list = [unet for unet in self.ema_unets]
            delattr(self, 'ema_unets')
            self.ema_unets = unets_list

        if index != self.ema_unet_being_trained_index:
            for unet_index, unet in enumerate(self.ema_unets):
                unet.to(self.device if unet_index == index else 'cpu')

        self.ema_unet_being_trained_index = index
        return self.ema_unets[index]


Prompt:
ation dataloader was already added'

        dl = DataLoader(ds, batch_size = batch_size, **dl_kwargs)
        self.add_valid_dataloader(dl)

    def create_train_iter(self):
        assert exists(self.train_dl), 'training dataloader has not been registered with the trainer yet'

        if exists(self.train_dl_iter):
            return

        self.train_dl_iter = cycle(self.train_dl)

    def create_valid_iter(self):
        assert exists(self.valid_dl), 'validation dataloader has not been registered with the trainer yet'

        if exists(self.valid_dl_iter):
            return

        self.valid_dl_iter = cycle(self.valid_dl)

    def train_step(self, unet_number = None, **kwargs):
        if not self.prepared:
            self.prepare()
        self.create_train_iter()
        loss = self.step_with_dl_iter(self.train_dl_iter, unet_number = unet_number, **kwargs)
        self.update(unet_number = unet_number)
        return loss

    @torch.no_grad()
    @eval_decorator
    def valid_step(self, **kwargs):
        if not self.prepared:
            self.prepare()
        self.create_valid_iter()
        context = self.use_ema_unets if kwargs.pop('use_ema_unets', False) else nullcontext
        with context():
            loss = self.step_with_dl_iter(self.valid_dl_iter, **kwargs)
        return loss

    def step_with_dl_iter(self, dl_iter, **kwargs):
        dl_tuple_output = cast_tuple(next(dl_iter))
        model_input = dict(list(zip(self.dl_tuple_output_keywords_names, dl_tuple_output)))
        loss = self.forward(**{**kwargs, **model_input})
        return loss

    # checkpointing functions

    @property
    def all_checkpoints_sorted(self):
        glob_pattern = os.path.join(self.checkpoint_path, '*.pt')
        checkpoints = self.fs.glob(glob_pattern)
        sorted_checkpoints = sorted(checkpoints, key = lambda x: int(str(x).split('.')[-2]), reverse = True)
        return sorted_checkpoints

    def load_from_checkpoint_folder(self, last_total_steps = -1):
        if last_total_steps != -1:
            filepath = os.path.join(self.checkpoint_path, f'checkpoint.{last_total_steps}.pt')
            self.load(filepath)
            return

        sorted_checkpoints = self.all_checkpoints_sorted

        if len(sorted_checkpoints) == 0:
            self.print(f'no checkpoints found to load from at {self.checkpoint_path}')
            return

        last_checkpoint = sorted_checkpoints[0]
        self.load(last_checkpoint)

    def save_to_checkpoint_folder(self):
        self.accelerator.wait_for_everyone()

        if not self.can_checkpoint:
            return

        total_steps = int(self.steps.sum().item())
        filepath = os.path.join(self.checkpoint_path, f'checkpoint.{total_steps}.pt')

        self.save(filepath)

        if self.max_checkpoints_keep <= 0:
            return

        sorted_checkpoints = self.all_checkpoints_sorted
        checkpoints_to_discard = sorted_checkpoints[self.max_checkpoints_keep:]

        for checkpoint in checkpoints_to_discard:
            self.fs.rm(checkpoint)

    # saving and loading functions

    def save(
        self,
        path,
        overwrite = True,
        without_optim_and_sched = False,
        **kwargs
    ):
        self.accelerator.wait_for_everyone()

        if not self.can_checkpoint:
            return

        fs = self.fs

        assert not (fs.exists(path) and not overwrite)

        self.reset_ema_unets_all_one_device()

        save_obj = dict(
            model = self.imagen.state_dict(),
            version = __version__,
            steps = self.steps.cpu(),
            **kwargs
        )

        save_optim_and_sched_iter = range(0, self.num_unets) if not without_optim_and_sched else tuple()

        for ind in save_optim_and_sched_iter:
            scaler_key = f'scaler{ind}'
            optimizer_key = f'optim{ind}'
            scheduler_key = f'scheduler{ind}'
            warmup_scheduler_key = f'warmup{ind}'

            scaler = getattr(self, scaler_key)
            optimizer = getattr(self, optimizer_key)
            scheduler = getattr(self, scheduler_key)
            warmup_scheduler = getattr(self, warmup_scheduler_key)

            if exists(scheduler):
                save_obj = {**save_obj, scheduler_key: scheduler.state_dict()}

            if exists(warmup_scheduler):
                save_obj = {**save_obj, warmup_scheduler_key: warmup_scheduler.state_dict()}

            save_obj = {**save_obj, scaler_key: scaler.state_dict(), optimizer_key: optimizer.state_dict()}

        if self.use_ema:
            save_obj = {**save_obj, 'ema': self.ema_unets.state_dict()}

        # determine if imagen config is available

        if hasattr(self.imagen, '_config'):
            self.print(f'this checkpoint is commandable from the CLI - "imagen --model {str(path)} \"<prompt>\""')

            save_obj = {
                **save_obj,
                'imagen_type': 'elucidated' if self.is_elucidated else 'original',
                'imagen_params': self.imagen._config
            }

        #save to path

        with fs.open(path, 'wb') as f:
            torch.save(save_obj, f)

        self.print(f'checkpoint saved to {path}')

    def load(self, path, only_model = False, strict = True, noop_if_not_exist = False):
        fs = self.fs

        if noop_if_not_exist and not fs.exists(path):
            self.print(f'trainer checkpoint not found at {str(path)}')
            return

        assert fs.exists(path), f'{path} does not exist'

        self.reset_ema_unets_all_one_device()

        # to avoid extra GPU memory usage in main process when using Accelerate

        with fs.open(path) as f:
            loaded_obj = torch.load(f, map_location='cpu')

        if version.parse(__version__) != version.parse(loaded_obj['version']):
            self.print(f'loading saved imagen at version {loaded_obj["version"]}, but current package version is {__version__}')

        try:
            self.imagen.load_state_dict(loaded_obj['model'], strict = strict)
        except RuntimeError:
            print("Failed loading state dict. Trying partial load")
            self.imagen.load_state_dict(restore_parts(self.imagen.state_dict(),
                                                      loaded_obj['model']))

        if only_model:
            return loaded_obj

        self.steps.copy_(loaded_obj['steps'])

        for ind in range(0, self.num_unets):
            scaler_key = f'scaler{ind}'
            optimizer_key = f'optim{ind}'
            scheduler_key = f'scheduler{ind}'
            warmup_scheduler_key = f'warmup{ind}'

            scaler = getattr(self, scaler_key)
            optimizer = getattr(self, optimizer_key)
            scheduler = getattr(self, scheduler_key)
            warmup_scheduler = getattr(self, warmup_scheduler_key)

            if exists(scheduler) and scheduler_key in loaded_obj:
                scheduler.load_state_dict(loaded_obj[scheduler_key])

            if exists(warmup_scheduler) and warmup_scheduler_key in loaded_obj:
                warmup_scheduler.load_state_dict(loaded_obj[warmup_scheduler_key])

            if exists(optimizer):
                try:
                    optimizer.load_state_dict(loaded_obj[optimizer_key])
                    scaler.load_state_dict(loaded_obj[scaler_key])
                except:
                    self.print('could not load optimizer and scaler, possibly because you have turned on mixed precision training since the last run. resuming with new optimizer and scalers')

        if self.use_ema:
            assert 'ema' in loaded_obj
            try:
                self.ema_unets.load_state_dict(loaded_obj['ema'], strict = strict)
            except RuntimeError:
                print("Failed loading state dict. Trying partial load")
                self.ema_unets.load_state_dict(restore_parts(self.ema_unets.state_dict(),
                                                             loaded_obj['ema']))

        self.print(f'checkpoint loaded from {path}')
        return loaded_obj

    # managing ema unets and their devices

    @property
    def unets(self):
        return nn.ModuleList([ema.ema_model for ema in self.ema_unets])

    def get_ema_unet(self, unet_number = None):
        if not self.use_ema:
            return

        unet_number = self.validate_unet_number(unet_number)
        index = unet_number - 1

        if isinstance(self.unets, nn.ModuleList):
            unets_list = [unet for unet in self.ema_unets]
            delattr(self, 'ema_unets')
            self.ema_unets = unets_list

        if index != self.ema_unet_being_trained_index:
            for unet_index, unet in enumerate(self.ema_unets):
                unet.to(self.device if unet_index == index else 'cpu')

        self.ema_unet_being_trained_index = index
        return self.ema_unets[index]

    def reset_ema_unets_all_one_device(self, device = None):
        if not self.use_ema:
            return

        device = default(device, self.device)
        self.ema_unets = nn.ModuleList([*self.ema_unets])
        self.ema_unets.to(device)

        self.ema_unet_being_trained_index = -1

    @torch.no_grad()
    @contextmanager
    def use_ema_unets(self):
        if not self.use_ema:
            output = yield
            return output

        self.reset_ema_unets_all_one_device()
        self.imagen.reset_unets_all_one_device()

        self.unets.eval()

        trainable_unets = self.imagen.unets
        self.imagen.unets = self.unets                  # swap in exponential moving averaged unets for sampling

        output = yield

        self.imagen.unets = trainable_unets             # restore original training unets

        # cast the ema_model unets back to original device
        for ema in self.ema_unets:
            ema.restore_ema_model_device()

        return output

    def print_unet_devices(self):
        self.print('unet devices:')
        for i, unet in enumerate(self.imagen.unets):
            device = next(unet.parameters()).device
            self.print(f'\tunet {i}: {device}')

        if not self.use_ema:
            return

        self.print('\nema unet devices:')
        for i, ema_unet in enumerate(self.ema_unets):
            device = next(ema_unet.parameters()).device
            self.print(f'\tema unet {i}: {device}')

    # overriding state dict functions

    def state_dict(self, *args, **kwargs):
        self.reset_ema_unets_all_one_device()
        return super().state_dict(*args, **kwargs)

    def load_state_dict(self, *args, **kwargs):
        self.reset_ema_unets_all_one_device()
        return super().load_state_dict(*args, **kwargs)

    # encoding text functions

    def encode_text(self, text, **kwargs):
        return self.imagen.encode_text(text, **kwargs)

    # forwarding functions and gradient step updates

    def update(self, unet_number = None):
        unet_number = self.validate_unet_number(unet_number)
        self.validate_and_set_unet_being_trained(unet_number)
        self.set_accelerator_scaler(unet_number)

        index = unet_number - 1
        unet = self.unet_being_trained

        optimizer = getattr(self, f'optim{index}')
        scaler = getattr(self, f'scaler{index}')
        scheduler = getattr(self, f'scheduler{index}')
        warmup_scheduler = getattr(self, f'warmup{index}')

        # set the grad scaler on the accelerator, since we are managing one per u-net

        if exists(self.max_grad_norm):
            self.accelerator.clip_grad_norm_(unet.parameters(), self.max_grad_norm)

        optimizer.step()
        optimizer.zero_grad()

        if self.use_ema:
            ema_unet = self.get_ema_unet(unet_number)
            ema_unet.update()

        # scheduler, if needed

        maybe_warmup_context = nullcontext() if not exists(warmup_scheduler) else warmup_scheduler.dampening()

        with maybe_warmup_context:
            if exists(scheduler) and not self.accelerator.optimizer_step_was_skipped: # recommended in the docs
                scheduler.step()

        self.steps += F.one_hot(torch.tensor(unet_number - 1, device = self.steps.device), num_classes = len(self.steps))

        if not exists(self.checkpoint_path):
            return

        total_steps = int(self.steps.sum().item())

        if total_steps % self.checkpoint_every:
            return

        self.save_to_checkpoint_folder()

    @torch.no_grad()
    @cast_torch_tensor
    @imagen_sample_in_chunks
    def sample(self, *args, **kwargs):
        context = nullcontext if  kwargs.pop('use_non_ema', False) else self.use_ema_unets

        self.print_untrained_unets()

        if not self.is_main:
            kwargs['use_tqdm'] = False

        with context():
            output = self.imagen.sample(*args, device = self.device, **kwargs)

        return output

    @partial(cast_torch_tensor, cast_fp16 = True)
    def forward(
        self,
        *args,
        unet_number = None,
        max_batch_size = None,
        **kwargs
    ):

        unet_number = self.validate_unet_number(unet_number)
        self.validate_and_set_unet_being_trained(unet_number)
        self.set_accelerator_scaler(unet_number)

        assert not exists(self.only_train_unet_number) or self.only_train_unet_number == unet_number, f'you can only train unet #{self.only_train_unet_number}'

        total_loss = 0.

        for chunk_size_frac, (chunked_args, chunked_kwargs) in split_args_and_kwargs(*args, split_size = max_batch_size, **kwargs):
            with self.accelerator.autocast():
                loss = self.imagen(*chunked_args, unet = self.unet_being_trained, unet_number = unet_number, **chunked_kwargs)
                loss = loss * chunk_size_frac

            total_loss += loss.item()

            if self.training:
                self.accelerator.backward(loss)

        return total_loss


Prompt:
# Copyright 2022 DeepMind Technologies Limited. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""RASP Evaluator which applies causal masks to selectors."""

from typing import Sequence, Union

import numpy as np
from tracr.rasp import rasp

class CausalEvaluator(rasp.DefaultRASPEvaluator):
  """Evaluates RASP with causal masking."""

  def evaluate(
      self, expr: rasp.RASPExpr, xs: Sequence[rasp.Value]
  ) -> Union[Sequence[rasp.Value], rasp.SelectorValue]:

    out = super().evaluate(expr, xs)

    if not isinstance(expr, rasp.Selector):
      return out

    out = np.array(out)
    causal_mask = np.tril(np.full(out.shape, 1))
    return np.logical_and(causal_mask, out).tolist()


Prompt:
# Copyright 2022 DeepMind Technologies Limited. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""RASP program objects.

Every object in the RASP language is a function.

The most important type is S-Op, which is a function List[Value] -> List[Value].

An S-Op represents a state inside the residual stream of the transformer.
Therefore, any RASP program that represents a transformer computation must
define a final S-Op that represents the state of the residual stream at the
end of the computation. In particular, given an S-Op `x`,
`x([1, 2, 3])` represents something like the state of the residual stream
at location `x` when the transformer is fed [1, 2, 3] as input.

A secondary (but still important) type is Selector, which is a function
List[Value] -> List[List[bool]]. Given a Selector `sel`, sel([1, 2, 3])
represents something like an attention matrix in the transformer.

For a full reference on RASP, see https://arxiv.org/abs/2106.06981.
"""

import abc
import collections.abc
import copy
import enum
import functools
import itertools
from typing import (Any, Callable, Dict, Generic, List, Mapping, Optional,
                    Sequence, TypeVar, Union)

from absl import logging
import numpy as np
from typing_extensions import Protocol

SelectorValue = List[List[bool]]
NumericValue = Union[int, float]
Value = Union[None, int, float, str, bool]
VT = TypeVar("VT", bound=Value)
RASPExprT = TypeVar("RASPExprT", bound="RASPExpr")
SOpT = TypeVar("SOpT", bound="SOp")
T = TypeVar("T")

_NAME_KEY = "name"
_ENCODING_KEY = "encoding"

# These are run on every expression when it's initialised.
# Add your own annotators to this dict to add custom default annotations.
#
# For example, DEFAULT_ANNOTATORS['foo'] will provide the default value for
# expr.annotations['foo]. The annotator will get called lazily the first time
# that key is accessed.
#
# See the `default_name` annotator for a full example.
DEFAULT_ANNOTATORS: Dict[str, "Annotator"] = {}

class Annotator(Protocol):

  def __call__(self, expr: "RASPExpr") -> Any:
    """What annotation to add to `expr`."""

class _Annotations(collections.abc.Mapping):
  """Holds the expression's annotations.

  It's immutable to the user, but will attempt to generate default values
  lazily when missing keys are requested.
  """

  def __init__(self, expr, **kwargs: Any):
    self._expr = expr
    self._inner_dict: Dict[str, Any] = {**kwargs}

  def __getitem__(self, key: str) -> Any:

    if key not in self._inner_dict:
      if key not in DEFAULT_ANNOTATORS:
        raise KeyError(
            f"No annotation exists for key '{key}'. "
            f"Available keys: {list(*self.keys(), *DEFAULT_ANNOTATORS.keys())}")
      self._inner_dict[key] = DEFAULT_ANNOTATORS[key](self._expr)

    return self._inner_dict[key]


Prompt:
# Copyright 2022 DeepMind Technologies Limited. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""RASP program objects.

Every object in the RASP language is a function.

The most important type is S-Op, which is a function List[Value] -> List[Value].

An S-Op represents a state inside the residual stream of the transformer.
Therefore, any RASP program that represents a transformer computation must
define a final S-Op that represents the state of the residual stream at the
end of the computation. In particular, given an S-Op `x`,
`x([1, 2, 3])` represents something like the state of the residual stream
at location `x` when the transformer is fed [1, 2, 3] as input.

A secondary (but still important) type is Selector, which is a function
List[Value] -> List[List[bool]]. Given a Selector `sel`, sel([1, 2, 3])
represents something like an attention matrix in the transformer.

For a full reference on RASP, see https://arxiv.org/abs/2106.06981.
"""

import abc
import collections.abc
import copy
import enum
import functools
import itertools
from typing import (Any, Callable, Dict, Generic, List, Mapping, Optional,
                    Sequence, TypeVar, Union)

from absl import logging
import numpy as np
from typing_extensions import Protocol

SelectorValue = List[List[bool]]
NumericValue = Union[int, float]
Value = Union[None, int, float, str, bool]
VT = TypeVar("VT", bound=Value)
RASPExprT = TypeVar("RASPExprT", bound="RASPExpr")
SOpT = TypeVar("SOpT", bound="SOp")
T = TypeVar("T")

_NAME_KEY = "name"
_ENCODING_KEY = "encoding"

# These are run on every expression when it's initialised.
# Add your own annotators to this dict to add custom default annotations.
#
# For example, DEFAULT_ANNOTATORS['foo'] will provide the default value for
# expr.annotations['foo]. The annotator will get called lazily the first time
# that key is accessed.
#
# See the `default_name` annotator for a full example.
DEFAULT_ANNOTATORS: Dict[str, "Annotator"] = {}

class Annotator(Protocol):

  def __call__(self, expr: "RASPExpr") -> Any:
    """What annotation to add to `expr`."""

class _Annotations(collections.abc.Mapping):
  """Holds the expression's annotations.

  It's immutable to the user, but will attempt to generate default values
  lazily when missing keys are requested.
  """

  def __init__(self, expr, **kwargs: Any):
    self._expr = expr
    self._inner_dict: Dict[str, Any] = {**kwargs}

  def __getitem__(self, key: str) -> Any:
    if key not in self._inner_dict:
      if key not in DEFAULT_ANNOTATORS:
        raise KeyError(
            f"No annotation exists for key '{key}'. "
            f"Available keys: {list(*self.keys(), *DEFAULT_ANNOTATORS.keys())}")
      self._inner_dict[key] = DEFAULT_ANNOTATORS[key](self._expr)

    return self._inner_dict[key]

  def __iter__(self):
    return iter(self._inner_dict)

  def __len__(self):
    return len(self._inner_dict)

class RASPExpr(abc.ABC):
  """A class distinguishing RASP expressions from other objects."""
  _ids = itertools.count(1)

  def __init__(self):
    self._annotations: Mapping[str, Any] = _Annotations(self)

  @abc.abstractmethod
  def __call__(self,
               xs: Sequence[Value]) -> Union[Sequence[Value], SelectorValue]:
    """Evaluates the RASPExpr using the standard evaluator."""

  @property
  def annotations(self) -> Mapping[str, Any]:
    """The annotations of this expression instance."""
    return self._annotations

  @annotations.setter
  def annotations(self, annotations: Mapping[str, Any]):
    self._annotations = _Annotations(self, **annotations)

  @property
  def name(self) -> str:
    """The name of this expression."""
    return self.annotations[_NAME_KEY]

  @property
  @abc.abstractmethod
  def children(self) -> Sequence["RASPExpr"]:
    """Direct dependencies of this expression."""

  @functools.cached_property
  def unique_id(self):
    """A unique id for every expression instance."""
    return next(self._ids)

  def copy(self: RASPExprT) -> RASPExprT:
    """Returns a shallow copy of this RASPExpr with a new ID."""
    return copy.copy(self)

  @property
  def label(self) -> str:
    return f"{self.name}_{self.unique_id}"

  def named(self: RASPExprT, name: str) -> RASPExprT:
    """Convenience method for adding a name."""
    return annotate(self, name=name)

  def annotated(self: RASPExprT, **annotations) -> RASPExprT:
    """Convenience method for adding annotations."""
    return annotate(self, **annotations)

def annotate(expr: RASPExprT, **annotations) -> RASPExprT:
  """Creates a new expr with added annotations."""

  new = expr.copy()
  # Note that new annotations will overwrite existing ones with matching keys.
  new.annotations = {**expr.annotations, **annotations}
  return new


Prompt:
# Copyright 2022 DeepMind Technologies Limited. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""RASP program objects.

Every object in the RASP language is a function.

The most important type is S-Op, which is a function List[Value] -> List[Value].

An S-Op represents a state inside the residual stream of the transformer.
Therefore, any RASP program that represents a transformer computation must
define a final S-Op that represents the state of the residual stream at the
end of the computation. In particular, given an S-Op `x`,
`x([1, 2, 3])` represents something like the state of the residual stream
at location `x` when the transformer is fed [1, 2, 3] as input.

A secondary (but still important) type is Selector, which is a function
List[Value] -> List[List[bool]]. Given a Selector `sel`, sel([1, 2, 3])
represents something like an attention matrix in the transformer.

For a full reference on RASP, see https://arxiv.org/abs/2106.06981.
"""

import abc
import collections.abc
import copy
import enum
import functools
import itertools
from typing import (Any, Callable, Dict, Generic, List, Mapping, Optional,
                    Sequence, TypeVar, Union)

from absl import logging
import numpy as np
from typing_extensions import Protocol

SelectorValue = List[List[bool]]
NumericValue = Union[int, float]
Value = Union[None, int, float, str, bool]
VT = TypeVar("VT", bound=Value)
RASPExprT = TypeVar("RASPExprT", bound="RASPExpr")
SOpT = TypeVar("SOpT", bound="SOp")
T = TypeVar("T")

_NAME_KEY = "name"
_ENCODING_KEY = "encoding"

# These are run on every expression when it's initialised.
# Add your own annotators to this dict to add custom default annotations.
#
# For example, DEFAULT_ANNOTATORS['foo'] will provide the default value for
# expr.annotations['foo]. The annotator will get called lazily the first time
# that key is accessed.
#
# See the `default_name` annotator for a full example.
DEFAULT_ANNOTATORS: Dict[str, "Annotator"] = {}

class Annotator(Protocol):

  def __call__(self, expr: "RASPExpr") -> Any:
    """What annotation to add to `expr`."""

class _Annotations(collections.abc.Mapping):
  """Holds the expression's annotations.

  It's immutable to the user, but will attempt to generate default values
  lazily when missing keys are requested.
  """

  def __init__(self, expr, **kwargs: Any):
    self._expr = expr
    self._inner_dict: Dict[str, Any] = {**kwargs}

  def __getitem__(self, key: str) -> Any:
    if key not in self._inner_dict:
      if key not in DEFAULT_ANNOTATORS:
        raise KeyError(
            f"No annotation exists for key '{key}'. "
            f"Available keys: {list(*self.keys(), *DEFAULT_ANNOTATORS.keys())}")
      self._inner_dict[key] = DEFAULT_ANNOTATORS[key](self._expr)

    return self._inner_dict[key]

  def __iter__(self):
    return iter(self._inner_dict)

  def __len__(self):
    return len(self._inner_dict)

class RASPExpr(abc.ABC):
  """A class distinguishing RASP expressions from other objects."""
  _ids = itertools.count(1)

  def __init__(self):
    self._annotations: Mapping[str, Any] = _Annotations(self)

  @abc.abstractmethod
  def __call__(self,
               xs: Sequence[Value]) -> Union[Sequence[Value], SelectorValue]:
    """Evaluates the RASPExpr using the standard evaluator."""

  @property
  def annotations(self) -> Mapping[str, Any]:
    """The annotations of this expression instance."""
    return self._annotations

  @annotations.setter
  def annotations(self, annotations: Mapping[str, Any]):
    self._annotations = _Annotations(self, **annotations)

  @property
  def name(self) -> str:
    """The name of this expression."""
    return self.annotations[_NAME_KEY]

  @property
  @abc.abstractmethod
  def children(self) -> Sequence["RASPExpr"]:
    """Direct dependencies of this expression."""

  @functools.cached_property
  def unique_id(self):
    """A unique id for every expression instance."""
    return next(self._ids)

  def copy(self: RASPExprT) -> RASPExprT:
    """Returns a shallow copy of this RASPExpr with a new ID."""
    return copy.copy(self)

  @property
  def label(self) -> str:
    return f"{self.name}_{self.unique_id}"

  def named(self: RASPExprT, name: str) -> RASPExprT:
    """Convenience method for adding a name."""
    return annotate(self, name=name)

  def annotated(self: RASPExprT, **annotations) -> RASPExprT:
    """Convenience method for adding annotations."""
    return annotate(self, **annotations)

def annotate(expr: RASPExprT, **annotations) -> RASPExprT:
  """Creates a new expr with added annotations."""
  new = expr.copy()
  # Note that new annotations will overwrite existing ones with matching keys.
  new.annotations = {**expr.annotations, **annotations}
  return new

### S-Ops.

class SOp(RASPExpr):
  """A Sequence Operation."""

  def __call__(self, xs: Sequence[Value]) -> Sequence[Value]:
    return evaluate(self, xs)  # pytype: disable=bad-return-type

  # Allow construction of SOps using numeric operators with constant values.
  # Note: if inheriting SOp by a dataclass, make sure to disable eq and order,
  # as they will override these.

  def __lt__(self, other: Value) -> "SOp":
    """self < other."""
    return Map(lambda x: x < other, self)

  def __le__(self, other: Value) -> "SOp":
    """self <= other."""
    return Map(lambda x: x <= other, self)

  def __eq__(self, other: Value) -> "SOp":
    """self == other."""
    return Map(lambda x: x == other, self)

  def __ne__(self, other: Value) -> "SOp":
    """self != other."""
    return Map(lambda x: x != other, self)

  def __gt__(self, other: Value) -> "SOp":
    """self > other."""
    return Map(lambda x: x > other, self)

  def __ge__(self, other: Value) -> "SOp":
    """self >= other."""
    return Map(lambda x: x >= other, self)

  def __add__(self, other: Union["SOp", Value]) -> "SOp":
    """self + other."""

    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x + y, self, other)
    return Map(lambda x: x + other, self)


Prompt:
# Copyright 2022 DeepMind Technologies Limited. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""RASP program objects.

Every object in the RASP language is a function.

The most important type is S-Op, which is a function List[Value] -> List[Value].

An S-Op represents a state inside the residual stream of the transformer.
Therefore, any RASP program that represents a transformer computation must
define a final S-Op that represents the state of the residual stream at the
end of the computation. In particular, given an S-Op `x`,
`x([1, 2, 3])` represents something like the state of the residual stream
at location `x` when the transformer is fed [1, 2, 3] as input.

A secondary (but still important) type is Selector, which is a function
List[Value] -> List[List[bool]]. Given a Selector `sel`, sel([1, 2, 3])
represents something like an attention matrix in the transformer.

For a full reference on RASP, see https://arxiv.org/abs/2106.06981.
"""

import abc
import collections.abc
import copy
import enum
import functools
import itertools
from typing import (Any, Callable, Dict, Generic, List, Mapping, Optional,
                    Sequence, TypeVar, Union)

from absl import logging
import numpy as np
from typing_extensions import Protocol

SelectorValue = List[List[bool]]
NumericValue = Union[int, float]
Value = Union[None, int, float, str, bool]
VT = TypeVar("VT", bound=Value)
RASPExprT = TypeVar("RASPExprT", bound="RASPExpr")
SOpT = TypeVar("SOpT", bound="SOp")
T = TypeVar("T")

_NAME_KEY = "name"
_ENCODING_KEY = "encoding"

# These are run on every expression when it's initialised.
# Add your own annotators to this dict to add custom default annotations.
#
# For example, DEFAULT_ANNOTATORS['foo'] will provide the default value for
# expr.annotations['foo]. The annotator will get called lazily the first time
# that key is accessed.
#
# See the `default_name` annotator for a full example.
DEFAULT_ANNOTATORS: Dict[str, "Annotator"] = {}

class Annotator(Protocol):

  def __call__(self, expr: "RASPExpr") -> Any:
    """What annotation to add to `expr`."""

class _Annotations(collections.abc.Mapping):
  """Holds the expression's annotations.

  It's immutable to the user, but will attempt to generate default values
  lazily when missing keys are requested.
  """

  def __init__(self, expr, **kwargs: Any):
    self._expr = expr
    self._inner_dict: Dict[str, Any] = {**kwargs}

  def __getitem__(self, key: str) -> Any:
    if key not in self._inner_dict:
      if key not in DEFAULT_ANNOTATORS:
        raise KeyError(
            f"No annotation exists for key '{key}'. "
            f"Available keys: {list(*self.keys(), *DEFAULT_ANNOTATORS.keys())}")
      self._inner_dict[key] = DEFAULT_ANNOTATORS[key](self._expr)

    return self._inner_dict[key]

  def __iter__(self):
    return iter(self._inner_dict)

  def __len__(self):
    return len(self._inner_dict)

class RASPExpr(abc.ABC):
  """A class distinguishing RASP expressions from other objects."""
  _ids = itertools.count(1)

  def __init__(self):
    self._annotations: Mapping[str, Any] = _Annotations(self)

  @abc.abstractmethod
  def __call__(self,
               xs: Sequence[Value]) -> Union[Sequence[Value], SelectorValue]:
    """Evaluates the RASPExpr using the standard evaluator."""

  @property
  def annotations(self) -> Mapping[str, Any]:
    """The annotations of this expression instance."""
    return self._annotations

  @annotations.setter
  def annotations(self, annotations: Mapping[str, Any]):
    self._annotations = _Annotations(self, **annotations)

  @property
  def name(self) -> str:
    """The name of this expression."""
    return self.annotations[_NAME_KEY]

  @property
  @abc.abstractmethod
  def children(self) -> Sequence["RASPExpr"]:
    """Direct dependencies of this expression."""

  @functools.cached_property
  def unique_id(self):
    """A unique id for every expression instance."""
    return next(self._ids)

  def copy(self: RASPExprT) -> RASPExprT:
    """Returns a shallow copy of this RASPExpr with a new ID."""
    return copy.copy(self)

  @property
  def label(self) -> str:
    return f"{self.name}_{self.unique_id}"

  def named(self: RASPExprT, name: str) -> RASPExprT:
    """Convenience method for adding a name."""
    return annotate(self, name=name)

  def annotated(self: RASPExprT, **annotations) -> RASPExprT:
    """Convenience method for adding annotations."""
    return annotate(self, **annotations)

def annotate(expr: RASPExprT, **annotations) -> RASPExprT:
  """Creates a new expr with added annotations."""
  new = expr.copy()
  # Note that new annotations will overwrite existing ones with matching keys.
  new.annotations = {**expr.annotations, **annotations}
  return new

### S-Ops.

class SOp(RASPExpr):
  """A Sequence Operation."""

  def __call__(self, xs: Sequence[Value]) -> Sequence[Value]:
    return evaluate(self, xs)  # pytype: disable=bad-return-type

  # Allow construction of SOps using numeric operators with constant values.
  # Note: if inheriting SOp by a dataclass, make sure to disable eq and order,
  # as they will override these.

  def __lt__(self, other: Value) -> "SOp":
    """self < other."""
    return Map(lambda x: x < other, self)

  def __le__(self, other: Value) -> "SOp":
    """self <= other."""
    return Map(lambda x: x <= other, self)

  def __eq__(self, other: Value) -> "SOp":
    """self == other."""
    return Map(lambda x: x == other, self)

  def __ne__(self, other: Value) -> "SOp":
    """self != other."""
    return Map(lambda x: x != other, self)

  def __gt__(self, other: Value) -> "SOp":
    """self > other."""
    return Map(lambda x: x > other, self)

  def __ge__(self, other: Value) -> "SOp":
    """self >= other."""
    return Map(lambda x: x >= other, self)

  def __add__(self, other: Union["SOp", Value]) -> "SOp":
    """self + other."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x + y, self, other)
    return Map(lambda x: x + other, self)

  def __radd__(self, other: Union["SOp", Value]) -> "SOp":
    """other + self."""

    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x + y, other, self)
    return Map(lambda x: other + x, self)


Prompt:
# Copyright 2022 DeepMind Technologies Limited. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""RASP program objects.

Every object in the RASP language is a function.

The most important type is S-Op, which is a function List[Value] -> List[Value].

An S-Op represents a state inside the residual stream of the transformer.
Therefore, any RASP program that represents a transformer computation must
define a final S-Op that represents the state of the residual stream at the
end of the computation. In particular, given an S-Op `x`,
`x([1, 2, 3])` represents something like the state of the residual stream
at location `x` when the transformer is fed [1, 2, 3] as input.

A secondary (but still important) type is Selector, which is a function
List[Value] -> List[List[bool]]. Given a Selector `sel`, sel([1, 2, 3])
represents something like an attention matrix in the transformer.

For a full reference on RASP, see https://arxiv.org/abs/2106.06981.
"""

import abc
import collections.abc
import copy
import enum
import functools
import itertools
from typing import (Any, Callable, Dict, Generic, List, Mapping, Optional,
                    Sequence, TypeVar, Union)

from absl import logging
import numpy as np
from typing_extensions import Protocol

SelectorValue = List[List[bool]]
NumericValue = Union[int, float]
Value = Union[None, int, float, str, bool]
VT = TypeVar("VT", bound=Value)
RASPExprT = TypeVar("RASPExprT", bound="RASPExpr")
SOpT = TypeVar("SOpT", bound="SOp")
T = TypeVar("T")

_NAME_KEY = "name"
_ENCODING_KEY = "encoding"

# These are run on every expression when it's initialised.
# Add your own annotators to this dict to add custom default annotations.
#
# For example, DEFAULT_ANNOTATORS['foo'] will provide the default value for
# expr.annotations['foo]. The annotator will get called lazily the first time
# that key is accessed.
#
# See the `default_name` annotator for a full example.
DEFAULT_ANNOTATORS: Dict[str, "Annotator"] = {}

class Annotator(Protocol):

  def __call__(self, expr: "RASPExpr") -> Any:
    """What annotation to add to `expr`."""

class _Annotations(collections.abc.Mapping):
  """Holds the expression's annotations.

  It's immutable to the user, but will attempt to generate default values
  lazily when missing keys are requested.
  """

  def __init__(self, expr, **kwargs: Any):
    self._expr = expr
    self._inner_dict: Dict[str, Any] = {**kwargs}

  def __getitem__(self, key: str) -> Any:
    if key not in self._inner_dict:
      if key not in DEFAULT_ANNOTATORS:
        raise KeyError(
            f"No annotation exists for key '{key}'. "
            f"Available keys: {list(*self.keys(), *DEFAULT_ANNOTATORS.keys())}")
      self._inner_dict[key] = DEFAULT_ANNOTATORS[key](self._expr)

    return self._inner_dict[key]

  def __iter__(self):
    return iter(self._inner_dict)

  def __len__(self):
    return len(self._inner_dict)

class RASPExpr(abc.ABC):
  """A class distinguishing RASP expressions from other objects."""
  _ids = itertools.count(1)

  def __init__(self):
    self._annotations: Mapping[str, Any] = _Annotations(self)

  @abc.abstractmethod
  def __call__(self,
               xs: Sequence[Value]) -> Union[Sequence[Value], SelectorValue]:
    """Evaluates the RASPExpr using the standard evaluator."""

  @property
  def annotations(self) -> Mapping[str, Any]:
    """The annotations of this expression instance."""
    return self._annotations

  @annotations.setter
  def annotations(self, annotations: Mapping[str, Any]):
    self._annotations = _Annotations(self, **annotations)

  @property
  def name(self) -> str:
    """The name of this expression."""
    return self.annotations[_NAME_KEY]

  @property
  @abc.abstractmethod
  def children(self) -> Sequence["RASPExpr"]:
    """Direct dependencies of this expression."""

  @functools.cached_property
  def unique_id(self):
    """A unique id for every expression instance."""
    return next(self._ids)

  def copy(self: RASPExprT) -> RASPExprT:
    """Returns a shallow copy of this RASPExpr with a new ID."""
    return copy.copy(self)

  @property
  def label(self) -> str:
    return f"{self.name}_{self.unique_id}"

  def named(self: RASPExprT, name: str) -> RASPExprT:
    """Convenience method for adding a name."""
    return annotate(self, name=name)

  def annotated(self: RASPExprT, **annotations) -> RASPExprT:
    """Convenience method for adding annotations."""
    return annotate(self, **annotations)

def annotate(expr: RASPExprT, **annotations) -> RASPExprT:
  """Creates a new expr with added annotations."""
  new = expr.copy()
  # Note that new annotations will overwrite existing ones with matching keys.
  new.annotations = {**expr.annotations, **annotations}
  return new

### S-Ops.

class SOp(RASPExpr):
  """A Sequence Operation."""

  def __call__(self, xs: Sequence[Value]) -> Sequence[Value]:
    return evaluate(self, xs)  # pytype: disable=bad-return-type

  # Allow construction of SOps using numeric operators with constant values.
  # Note: if inheriting SOp by a dataclass, make sure to disable eq and order,
  # as they will override these.

  def __lt__(self, other: Value) -> "SOp":
    """self < other."""
    return Map(lambda x: x < other, self)

  def __le__(self, other: Value) -> "SOp":
    """self <= other."""
    return Map(lambda x: x <= other, self)

  def __eq__(self, other: Value) -> "SOp":
    """self == other."""
    return Map(lambda x: x == other, self)

  def __ne__(self, other: Value) -> "SOp":
    """self != other."""
    return Map(lambda x: x != other, self)

  def __gt__(self, other: Value) -> "SOp":
    """self > other."""
    return Map(lambda x: x > other, self)

  def __ge__(self, other: Value) -> "SOp":
    """self >= other."""
    return Map(lambda x: x >= other, self)

  def __add__(self, other: Union["SOp", Value]) -> "SOp":
    """self + other."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x + y, self, other)
    return Map(lambda x: x + other, self)

  def __radd__(self, other: Union["SOp", Value]) -> "SOp":
    """other + self."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x + y, other, self)
    return Map(lambda x: other + x, self)

  def __sub__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """self - other."""

    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x - y, self, other)
    return Map(lambda x: x - other, self)


Prompt:
# Copyright 2022 DeepMind Technologies Limited. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""RASP program objects.

Every object in the RASP language is a function.

The most important type is S-Op, which is a function List[Value] -> List[Value].

An S-Op represents a state inside the residual stream of the transformer.
Therefore, any RASP program that represents a transformer computation must
define a final S-Op that represents the state of the residual stream at the
end of the computation. In particular, given an S-Op `x`,
`x([1, 2, 3])` represents something like the state of the residual stream
at location `x` when the transformer is fed [1, 2, 3] as input.

A secondary (but still important) type is Selector, which is a function
List[Value] -> List[List[bool]]. Given a Selector `sel`, sel([1, 2, 3])
represents something like an attention matrix in the transformer.

For a full reference on RASP, see https://arxiv.org/abs/2106.06981.
"""

import abc
import collections.abc
import copy
import enum
import functools
import itertools
from typing import (Any, Callable, Dict, Generic, List, Mapping, Optional,
                    Sequence, TypeVar, Union)

from absl import logging
import numpy as np
from typing_extensions import Protocol

SelectorValue = List[List[bool]]
NumericValue = Union[int, float]
Value = Union[None, int, float, str, bool]
VT = TypeVar("VT", bound=Value)
RASPExprT = TypeVar("RASPExprT", bound="RASPExpr")
SOpT = TypeVar("SOpT", bound="SOp")
T = TypeVar("T")

_NAME_KEY = "name"
_ENCODING_KEY = "encoding"

# These are run on every expression when it's initialised.
# Add your own annotators to this dict to add custom default annotations.
#
# For example, DEFAULT_ANNOTATORS['foo'] will provide the default value for
# expr.annotations['foo]. The annotator will get called lazily the first time
# that key is accessed.
#
# See the `default_name` annotator for a full example.
DEFAULT_ANNOTATORS: Dict[str, "Annotator"] = {}

class Annotator(Protocol):

  def __call__(self, expr: "RASPExpr") -> Any:
    """What annotation to add to `expr`."""

class _Annotations(collections.abc.Mapping):
  """Holds the expression's annotations.

  It's immutable to the user, but will attempt to generate default values
  lazily when missing keys are requested.
  """

  def __init__(self, expr, **kwargs: Any):
    self._expr = expr
    self._inner_dict: Dict[str, Any] = {**kwargs}

  def __getitem__(self, key: str) -> Any:
    if key not in self._inner_dict:
      if key not in DEFAULT_ANNOTATORS:
        raise KeyError(
            f"No annotation exists for key '{key}'. "
            f"Available keys: {list(*self.keys(), *DEFAULT_ANNOTATORS.keys())}")
      self._inner_dict[key] = DEFAULT_ANNOTATORS[key](self._expr)

    return self._inner_dict[key]

  def __iter__(self):
    return iter(self._inner_dict)

  def __len__(self):
    return len(self._inner_dict)

class RASPExpr(abc.ABC):
  """A class distinguishing RASP expressions from other objects."""
  _ids = itertools.count(1)

  def __init__(self):
    self._annotations: Mapping[str, Any] = _Annotations(self)

  @abc.abstractmethod
  def __call__(self,
               xs: Sequence[Value]) -> Union[Sequence[Value], SelectorValue]:
    """Evaluates the RASPExpr using the standard evaluator."""

  @property
  def annotations(self) -> Mapping[str, Any]:
    """The annotations of this expression instance."""
    return self._annotations

  @annotations.setter
  def annotations(self, annotations: Mapping[str, Any]):
    self._annotations = _Annotations(self, **annotations)

  @property
  def name(self) -> str:
    """The name of this expression."""
    return self.annotations[_NAME_KEY]

  @property
  @abc.abstractmethod
  def children(self) -> Sequence["RASPExpr"]:
    """Direct dependencies of this expression."""

  @functools.cached_property
  def unique_id(self):
    """A unique id for every expression instance."""
    return next(self._ids)

  def copy(self: RASPExprT) -> RASPExprT:
    """Returns a shallow copy of this RASPExpr with a new ID."""
    return copy.copy(self)

  @property
  def label(self) -> str:
    return f"{self.name}_{self.unique_id}"

  def named(self: RASPExprT, name: str) -> RASPExprT:
    """Convenience method for adding a name."""
    return annotate(self, name=name)

  def annotated(self: RASPExprT, **annotations) -> RASPExprT:
    """Convenience method for adding annotations."""
    return annotate(self, **annotations)

def annotate(expr: RASPExprT, **annotations) -> RASPExprT:
  """Creates a new expr with added annotations."""
  new = expr.copy()
  # Note that new annotations will overwrite existing ones with matching keys.
  new.annotations = {**expr.annotations, **annotations}
  return new

### S-Ops.

class SOp(RASPExpr):
  """A Sequence Operation."""

  def __call__(self, xs: Sequence[Value]) -> Sequence[Value]:
    return evaluate(self, xs)  # pytype: disable=bad-return-type

  # Allow construction of SOps using numeric operators with constant values.
  # Note: if inheriting SOp by a dataclass, make sure to disable eq and order,
  # as they will override these.

  def __lt__(self, other: Value) -> "SOp":
    """self < other."""
    return Map(lambda x: x < other, self)

  def __le__(self, other: Value) -> "SOp":
    """self <= other."""
    return Map(lambda x: x <= other, self)

  def __eq__(self, other: Value) -> "SOp":
    """self == other."""
    return Map(lambda x: x == other, self)

  def __ne__(self, other: Value) -> "SOp":
    """self != other."""
    return Map(lambda x: x != other, self)

  def __gt__(self, other: Value) -> "SOp":
    """self > other."""
    return Map(lambda x: x > other, self)

  def __ge__(self, other: Value) -> "SOp":
    """self >= other."""
    return Map(lambda x: x >= other, self)

  def __add__(self, other: Union["SOp", Value]) -> "SOp":
    """self + other."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x + y, self, other)
    return Map(lambda x: x + other, self)

  def __radd__(self, other: Union["SOp", Value]) -> "SOp":
    """other + self."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x + y, other, self)
    return Map(lambda x: other + x, self)

  def __sub__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """self - other."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x - y, self, other)
    return Map(lambda x: x - other, self)

  def __rsub__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """other - self."""

    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x - y, other, self)
    return Map(lambda x: other - x, self)


Prompt:
# Copyright 2022 DeepMind Technologies Limited. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""RASP program objects.

Every object in the RASP language is a function.

The most important type is S-Op, which is a function List[Value] -> List[Value].

An S-Op represents a state inside the residual stream of the transformer.
Therefore, any RASP program that represents a transformer computation must
define a final S-Op that represents the state of the residual stream at the
end of the computation. In particular, given an S-Op `x`,
`x([1, 2, 3])` represents something like the state of the residual stream
at location `x` when the transformer is fed [1, 2, 3] as input.

A secondary (but still important) type is Selector, which is a function
List[Value] -> List[List[bool]]. Given a Selector `sel`, sel([1, 2, 3])
represents something like an attention matrix in the transformer.

For a full reference on RASP, see https://arxiv.org/abs/2106.06981.
"""

import abc
import collections.abc
import copy
import enum
import functools
import itertools
from typing import (Any, Callable, Dict, Generic, List, Mapping, Optional,
                    Sequence, TypeVar, Union)

from absl import logging
import numpy as np
from typing_extensions import Protocol

SelectorValue = List[List[bool]]
NumericValue = Union[int, float]
Value = Union[None, int, float, str, bool]
VT = TypeVar("VT", bound=Value)
RASPExprT = TypeVar("RASPExprT", bound="RASPExpr")
SOpT = TypeVar("SOpT", bound="SOp")
T = TypeVar("T")

_NAME_KEY = "name"
_ENCODING_KEY = "encoding"

# These are run on every expression when it's initialised.
# Add your own annotators to this dict to add custom default annotations.
#
# For example, DEFAULT_ANNOTATORS['foo'] will provide the default value for
# expr.annotations['foo]. The annotator will get called lazily the first time
# that key is accessed.
#
# See the `default_name` annotator for a full example.
DEFAULT_ANNOTATORS: Dict[str, "Annotator"] = {}

class Annotator(Protocol):

  def __call__(self, expr: "RASPExpr") -> Any:
    """What annotation to add to `expr`."""

class _Annotations(collections.abc.Mapping):
  """Holds the expression's annotations.

  It's immutable to the user, but will attempt to generate default values
  lazily when missing keys are requested.
  """

  def __init__(self, expr, **kwargs: Any):
    self._expr = expr
    self._inner_dict: Dict[str, Any] = {**kwargs}

  def __getitem__(self, key: str) -> Any:
    if key not in self._inner_dict:
      if key not in DEFAULT_ANNOTATORS:
        raise KeyError(
            f"No annotation exists for key '{key}'. "
            f"Available keys: {list(*self.keys(), *DEFAULT_ANNOTATORS.keys())}")
      self._inner_dict[key] = DEFAULT_ANNOTATORS[key](self._expr)

    return self._inner_dict[key]

  def __iter__(self):
    return iter(self._inner_dict)

  def __len__(self):
    return len(self._inner_dict)

class RASPExpr(abc.ABC):
  """A class distinguishing RASP expressions from other objects."""
  _ids = itertools.count(1)

  def __init__(self):
    self._annotations: Mapping[str, Any] = _Annotations(self)

  @abc.abstractmethod
  def __call__(self,
               xs: Sequence[Value]) -> Union[Sequence[Value], SelectorValue]:
    """Evaluates the RASPExpr using the standard evaluator."""

  @property
  def annotations(self) -> Mapping[str, Any]:
    """The annotations of this expression instance."""
    return self._annotations

  @annotations.setter
  def annotations(self, annotations: Mapping[str, Any]):
    self._annotations = _Annotations(self, **annotations)

  @property
  def name(self) -> str:
    """The name of this expression."""
    return self.annotations[_NAME_KEY]

  @property
  @abc.abstractmethod
  def children(self) -> Sequence["RASPExpr"]:
    """Direct dependencies of this expression."""

  @functools.cached_property
  def unique_id(self):
    """A unique id for every expression instance."""
    return next(self._ids)

  def copy(self: RASPExprT) -> RASPExprT:
    """Returns a shallow copy of this RASPExpr with a new ID."""
    return copy.copy(self)

  @property
  def label(self) -> str:
    return f"{self.name}_{self.unique_id}"

  def named(self: RASPExprT, name: str) -> RASPExprT:
    """Convenience method for adding a name."""
    return annotate(self, name=name)

  def annotated(self: RASPExprT, **annotations) -> RASPExprT:
    """Convenience method for adding annotations."""
    return annotate(self, **annotations)

def annotate(expr: RASPExprT, **annotations) -> RASPExprT:
  """Creates a new expr with added annotations."""
  new = expr.copy()
  # Note that new annotations will overwrite existing ones with matching keys.
  new.annotations = {**expr.annotations, **annotations}
  return new

### S-Ops.

class SOp(RASPExpr):
  """A Sequence Operation."""

  def __call__(self, xs: Sequence[Value]) -> Sequence[Value]:
    return evaluate(self, xs)  # pytype: disable=bad-return-type

  # Allow construction of SOps using numeric operators with constant values.
  # Note: if inheriting SOp by a dataclass, make sure to disable eq and order,
  # as they will override these.

  def __lt__(self, other: Value) -> "SOp":
    """self < other."""
    return Map(lambda x: x < other, self)

  def __le__(self, other: Value) -> "SOp":
    """self <= other."""
    return Map(lambda x: x <= other, self)

  def __eq__(self, other: Value) -> "SOp":
    """self == other."""
    return Map(lambda x: x == other, self)

  def __ne__(self, other: Value) -> "SOp":
    """self != other."""
    return Map(lambda x: x != other, self)

  def __gt__(self, other: Value) -> "SOp":
    """self > other."""
    return Map(lambda x: x > other, self)

  def __ge__(self, other: Value) -> "SOp":
    """self >= other."""
    return Map(lambda x: x >= other, self)

  def __add__(self, other: Union["SOp", Value]) -> "SOp":
    """self + other."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x + y, self, other)
    return Map(lambda x: x + other, self)

  def __radd__(self, other: Union["SOp", Value]) -> "SOp":
    """other + self."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x + y, other, self)
    return Map(lambda x: other + x, self)

  def __sub__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """self - other."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x - y, self, other)
    return Map(lambda x: x - other, self)

  def __rsub__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """other - self."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x - y, other, self)
    return Map(lambda x: other - x, self)

  def __mul__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """self * other."""

    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x * y, self, other)
    return Map(lambda x: x * other, self)


Prompt:
# Copyright 2022 DeepMind Technologies Limited. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""RASP program objects.

Every object in the RASP language is a function.

The most important type is S-Op, which is a function List[Value] -> List[Value].

An S-Op represents a state inside the residual stream of the transformer.
Therefore, any RASP program that represents a transformer computation must
define a final S-Op that represents the state of the residual stream at the
end of the computation. In particular, given an S-Op `x`,
`x([1, 2, 3])` represents something like the state of the residual stream
at location `x` when the transformer is fed [1, 2, 3] as input.

A secondary (but still important) type is Selector, which is a function
List[Value] -> List[List[bool]]. Given a Selector `sel`, sel([1, 2, 3])
represents something like an attention matrix in the transformer.

For a full reference on RASP, see https://arxiv.org/abs/2106.06981.
"""

import abc
import collections.abc
import copy
import enum
import functools
import itertools
from typing import (Any, Callable, Dict, Generic, List, Mapping, Optional,
                    Sequence, TypeVar, Union)

from absl import logging
import numpy as np
from typing_extensions import Protocol

SelectorValue = List[List[bool]]
NumericValue = Union[int, float]
Value = Union[None, int, float, str, bool]
VT = TypeVar("VT", bound=Value)
RASPExprT = TypeVar("RASPExprT", bound="RASPExpr")
SOpT = TypeVar("SOpT", bound="SOp")
T = TypeVar("T")

_NAME_KEY = "name"
_ENCODING_KEY = "encoding"

# These are run on every expression when it's initialised.
# Add your own annotators to this dict to add custom default annotations.
#
# For example, DEFAULT_ANNOTATORS['foo'] will provide the default value for
# expr.annotations['foo]. The annotator will get called lazily the first time
# that key is accessed.
#
# See the `default_name` annotator for a full example.
DEFAULT_ANNOTATORS: Dict[str, "Annotator"] = {}

class Annotator(Protocol):

  def __call__(self, expr: "RASPExpr") -> Any:
    """What annotation to add to `expr`."""

class _Annotations(collections.abc.Mapping):
  """Holds the expression's annotations.

  It's immutable to the user, but will attempt to generate default values
  lazily when missing keys are requested.
  """

  def __init__(self, expr, **kwargs: Any):
    self._expr = expr
    self._inner_dict: Dict[str, Any] = {**kwargs}

  def __getitem__(self, key: str) -> Any:
    if key not in self._inner_dict:
      if key not in DEFAULT_ANNOTATORS:
        raise KeyError(
            f"No annotation exists for key '{key}'. "
            f"Available keys: {list(*self.keys(), *DEFAULT_ANNOTATORS.keys())}")
      self._inner_dict[key] = DEFAULT_ANNOTATORS[key](self._expr)

    return self._inner_dict[key]

  def __iter__(self):
    return iter(self._inner_dict)

  def __len__(self):
    return len(self._inner_dict)

class RASPExpr(abc.ABC):
  """A class distinguishing RASP expressions from other objects."""
  _ids = itertools.count(1)

  def __init__(self):
    self._annotations: Mapping[str, Any] = _Annotations(self)

  @abc.abstractmethod
  def __call__(self,
               xs: Sequence[Value]) -> Union[Sequence[Value], SelectorValue]:
    """Evaluates the RASPExpr using the standard evaluator."""

  @property
  def annotations(self) -> Mapping[str, Any]:
    """The annotations of this expression instance."""
    return self._annotations

  @annotations.setter
  def annotations(self, annotations: Mapping[str, Any]):
    self._annotations = _Annotations(self, **annotations)

  @property
  def name(self) -> str:
    """The name of this expression."""
    return self.annotations[_NAME_KEY]

  @property
  @abc.abstractmethod
  def children(self) -> Sequence["RASPExpr"]:
    """Direct dependencies of this expression."""

  @functools.cached_property
  def unique_id(self):
    """A unique id for every expression instance."""
    return next(self._ids)

  def copy(self: RASPExprT) -> RASPExprT:
    """Returns a shallow copy of this RASPExpr with a new ID."""
    return copy.copy(self)

  @property
  def label(self) -> str:
    return f"{self.name}_{self.unique_id}"

  def named(self: RASPExprT, name: str) -> RASPExprT:
    """Convenience method for adding a name."""
    return annotate(self, name=name)

  def annotated(self: RASPExprT, **annotations) -> RASPExprT:
    """Convenience method for adding annotations."""
    return annotate(self, **annotations)

def annotate(expr: RASPExprT, **annotations) -> RASPExprT:
  """Creates a new expr with added annotations."""
  new = expr.copy()
  # Note that new annotations will overwrite existing ones with matching keys.
  new.annotations = {**expr.annotations, **annotations}
  return new

### S-Ops.

class SOp(RASPExpr):
  """A Sequence Operation."""

  def __call__(self, xs: Sequence[Value]) -> Sequence[Value]:
    return evaluate(self, xs)  # pytype: disable=bad-return-type

  # Allow construction of SOps using numeric operators with constant values.
  # Note: if inheriting SOp by a dataclass, make sure to disable eq and order,
  # as they will override these.

  def __lt__(self, other: Value) -> "SOp":
    """self < other."""
    return Map(lambda x: x < other, self)

  def __le__(self, other: Value) -> "SOp":
    """self <= other."""
    return Map(lambda x: x <= other, self)

  def __eq__(self, other: Value) -> "SOp":
    """self == other."""
    return Map(lambda x: x == other, self)

  def __ne__(self, other: Value) -> "SOp":
    """self != other."""
    return Map(lambda x: x != other, self)

  def __gt__(self, other: Value) -> "SOp":
    """self > other."""
    return Map(lambda x: x > other, self)

  def __ge__(self, other: Value) -> "SOp":
    """self >= other."""
    return Map(lambda x: x >= other, self)

  def __add__(self, other: Union["SOp", Value]) -> "SOp":
    """self + other."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x + y, self, other)
    return Map(lambda x: x + other, self)

  def __radd__(self, other: Union["SOp", Value]) -> "SOp":
    """other + self."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x + y, other, self)
    return Map(lambda x: other + x, self)

  def __sub__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """self - other."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x - y, self, other)
    return Map(lambda x: x - other, self)

  def __rsub__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """other - self."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x - y, other, self)
    return Map(lambda x: other - x, self)

  def __mul__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """self * other."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x * y, self, other)
    return Map(lambda x: x * other, self)

  def __rmul__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """other * self."""

    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x * y, other, self)
    return Map(lambda x: other * x, self)


Prompt:
# Copyright 2022 DeepMind Technologies Limited. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""RASP program objects.

Every object in the RASP language is a function.

The most important type is S-Op, which is a function List[Value] -> List[Value].

An S-Op represents a state inside the residual stream of the transformer.
Therefore, any RASP program that represents a transformer computation must
define a final S-Op that represents the state of the residual stream at the
end of the computation. In particular, given an S-Op `x`,
`x([1, 2, 3])` represents something like the state of the residual stream
at location `x` when the transformer is fed [1, 2, 3] as input.

A secondary (but still important) type is Selector, which is a function
List[Value] -> List[List[bool]]. Given a Selector `sel`, sel([1, 2, 3])
represents something like an attention matrix in the transformer.

For a full reference on RASP, see https://arxiv.org/abs/2106.06981.
"""

import abc
import collections.abc
import copy
import enum
import functools
import itertools
from typing import (Any, Callable, Dict, Generic, List, Mapping, Optional,
                    Sequence, TypeVar, Union)

from absl import logging
import numpy as np
from typing_extensions import Protocol

SelectorValue = List[List[bool]]
NumericValue = Union[int, float]
Value = Union[None, int, float, str, bool]
VT = TypeVar("VT", bound=Value)
RASPExprT = TypeVar("RASPExprT", bound="RASPExpr")
SOpT = TypeVar("SOpT", bound="SOp")
T = TypeVar("T")

_NAME_KEY = "name"
_ENCODING_KEY = "encoding"

# These are run on every expression when it's initialised.
# Add your own annotators to this dict to add custom default annotations.
#
# For example, DEFAULT_ANNOTATORS['foo'] will provide the default value for
# expr.annotations['foo]. The annotator will get called lazily the first time
# that key is accessed.
#
# See the `default_name` annotator for a full example.
DEFAULT_ANNOTATORS: Dict[str, "Annotator"] = {}

class Annotator(Protocol):

  def __call__(self, expr: "RASPExpr") -> Any:
    """What annotation to add to `expr`."""

class _Annotations(collections.abc.Mapping):
  """Holds the expression's annotations.

  It's immutable to the user, but will attempt to generate default values
  lazily when missing keys are requested.
  """

  def __init__(self, expr, **kwargs: Any):
    self._expr = expr
    self._inner_dict: Dict[str, Any] = {**kwargs}

  def __getitem__(self, key: str) -> Any:
    if key not in self._inner_dict:
      if key not in DEFAULT_ANNOTATORS:
        raise KeyError(
            f"No annotation exists for key '{key}'. "
            f"Available keys: {list(*self.keys(), *DEFAULT_ANNOTATORS.keys())}")
      self._inner_dict[key] = DEFAULT_ANNOTATORS[key](self._expr)

    return self._inner_dict[key]

  def __iter__(self):
    return iter(self._inner_dict)

  def __len__(self):
    return len(self._inner_dict)

class RASPExpr(abc.ABC):
  """A class distinguishing RASP expressions from other objects."""
  _ids = itertools.count(1)

  def __init__(self):
    self._annotations: Mapping[str, Any] = _Annotations(self)

  @abc.abstractmethod
  def __call__(self,
               xs: Sequence[Value]) -> Union[Sequence[Value], SelectorValue]:
    """Evaluates the RASPExpr using the standard evaluator."""

  @property
  def annotations(self) -> Mapping[str, Any]:
    """The annotations of this expression instance."""
    return self._annotations

  @annotations.setter
  def annotations(self, annotations: Mapping[str, Any]):
    self._annotations = _Annotations(self, **annotations)

  @property
  def name(self) -> str:
    """The name of this expression."""
    return self.annotations[_NAME_KEY]

  @property
  @abc.abstractmethod
  def children(self) -> Sequence["RASPExpr"]:
    """Direct dependencies of this expression."""

  @functools.cached_property
  def unique_id(self):
    """A unique id for every expression instance."""
    return next(self._ids)

  def copy(self: RASPExprT) -> RASPExprT:
    """Returns a shallow copy of this RASPExpr with a new ID."""
    return copy.copy(self)

  @property
  def label(self) -> str:
    return f"{self.name}_{self.unique_id}"

  def named(self: RASPExprT, name: str) -> RASPExprT:
    """Convenience method for adding a name."""
    return annotate(self, name=name)

  def annotated(self: RASPExprT, **annotations) -> RASPExprT:
    """Convenience method for adding annotations."""
    return annotate(self, **annotations)

def annotate(expr: RASPExprT, **annotations) -> RASPExprT:
  """Creates a new expr with added annotations."""
  new = expr.copy()
  # Note that new annotations will overwrite existing ones with matching keys.
  new.annotations = {**expr.annotations, **annotations}
  return new

### S-Ops.

class SOp(RASPExpr):
  """A Sequence Operation."""

  def __call__(self, xs: Sequence[Value]) -> Sequence[Value]:
    return evaluate(self, xs)  # pytype: disable=bad-return-type

  # Allow construction of SOps using numeric operators with constant values.
  # Note: if inheriting SOp by a dataclass, make sure to disable eq and order,
  # as they will override these.

  def __lt__(self, other: Value) -> "SOp":
    """self < other."""
    return Map(lambda x: x < other, self)

  def __le__(self, other: Value) -> "SOp":
    """self <= other."""
    return Map(lambda x: x <= other, self)

  def __eq__(self, other: Value) -> "SOp":
    """self == other."""
    return Map(lambda x: x == other, self)

  def __ne__(self, other: Value) -> "SOp":
    """self != other."""
    return Map(lambda x: x != other, self)

  def __gt__(self, other: Value) -> "SOp":
    """self > other."""
    return Map(lambda x: x > other, self)

  def __ge__(self, other: Value) -> "SOp":
    """self >= other."""
    return Map(lambda x: x >= other, self)

  def __add__(self, other: Union["SOp", Value]) -> "SOp":
    """self + other."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x + y, self, other)
    return Map(lambda x: x + other, self)

  def __radd__(self, other: Union["SOp", Value]) -> "SOp":
    """other + self."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x + y, other, self)
    return Map(lambda x: other + x, self)

  def __sub__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """self - other."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x - y, self, other)
    return Map(lambda x: x - other, self)

  def __rsub__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """other - self."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x - y, other, self)
    return Map(lambda x: other - x, self)

  def __mul__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """self * other."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x * y, self, other)
    return Map(lambda x: x * other, self)

  def __rmul__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """other * self."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x * y, other, self)
    return Map(lambda x: other * x, self)

  def __truediv__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """self / other."""

    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x / y, self, other)
    return Map(lambda x: x / other, self)


Prompt:
# Copyright 2022 DeepMind Technologies Limited. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""RASP program objects.

Every object in the RASP language is a function.

The most important type is S-Op, which is a function List[Value] -> List[Value].

An S-Op represents a state inside the residual stream of the transformer.
Therefore, any RASP program that represents a transformer computation must
define a final S-Op that represents the state of the residual stream at the
end of the computation. In particular, given an S-Op `x`,
`x([1, 2, 3])` represents something like the state of the residual stream
at location `x` when the transformer is fed [1, 2, 3] as input.

A secondary (but still important) type is Selector, which is a function
List[Value] -> List[List[bool]]. Given a Selector `sel`, sel([1, 2, 3])
represents something like an attention matrix in the transformer.

For a full reference on RASP, see https://arxiv.org/abs/2106.06981.
"""

import abc
import collections.abc
import copy
import enum
import functools
import itertools
from typing import (Any, Callable, Dict, Generic, List, Mapping, Optional,
                    Sequence, TypeVar, Union)

from absl import logging
import numpy as np
from typing_extensions import Protocol

SelectorValue = List[List[bool]]
NumericValue = Union[int, float]
Value = Union[None, int, float, str, bool]
VT = TypeVar("VT", bound=Value)
RASPExprT = TypeVar("RASPExprT", bound="RASPExpr")
SOpT = TypeVar("SOpT", bound="SOp")
T = TypeVar("T")

_NAME_KEY = "name"
_ENCODING_KEY = "encoding"

# These are run on every expression when it's initialised.
# Add your own annotators to this dict to add custom default annotations.
#
# For example, DEFAULT_ANNOTATORS['foo'] will provide the default value for
# expr.annotations['foo]. The annotator will get called lazily the first time
# that key is accessed.
#
# See the `default_name` annotator for a full example.
DEFAULT_ANNOTATORS: Dict[str, "Annotator"] = {}

class Annotator(Protocol):

  def __call__(self, expr: "RASPExpr") -> Any:
    """What annotation to add to `expr`."""

class _Annotations(collections.abc.Mapping):
  """Holds the expression's annotations.

  It's immutable to the user, but will attempt to generate default values
  lazily when missing keys are requested.
  """

  def __init__(self, expr, **kwargs: Any):
    self._expr = expr
    self._inner_dict: Dict[str, Any] = {**kwargs}

  def __getitem__(self, key: str) -> Any:
    if key not in self._inner_dict:
      if key not in DEFAULT_ANNOTATORS:
        raise KeyError(
            f"No annotation exists for key '{key}'. "
            f"Available keys: {list(*self.keys(), *DEFAULT_ANNOTATORS.keys())}")
      self._inner_dict[key] = DEFAULT_ANNOTATORS[key](self._expr)

    return self._inner_dict[key]

  def __iter__(self):
    return iter(self._inner_dict)

  def __len__(self):
    return len(self._inner_dict)

class RASPExpr(abc.ABC):
  """A class distinguishing RASP expressions from other objects."""
  _ids = itertools.count(1)

  def __init__(self):
    self._annotations: Mapping[str, Any] = _Annotations(self)

  @abc.abstractmethod
  def __call__(self,
               xs: Sequence[Value]) -> Union[Sequence[Value], SelectorValue]:
    """Evaluates the RASPExpr using the standard evaluator."""

  @property
  def annotations(self) -> Mapping[str, Any]:
    """The annotations of this expression instance."""
    return self._annotations

  @annotations.setter
  def annotations(self, annotations: Mapping[str, Any]):
    self._annotations = _Annotations(self, **annotations)

  @property
  def name(self) -> str:
    """The name of this expression."""
    return self.annotations[_NAME_KEY]

  @property
  @abc.abstractmethod
  def children(self) -> Sequence["RASPExpr"]:
    """Direct dependencies of this expression."""

  @functools.cached_property
  def unique_id(self):
    """A unique id for every expression instance."""
    return next(self._ids)

  def copy(self: RASPExprT) -> RASPExprT:
    """Returns a shallow copy of this RASPExpr with a new ID."""
    return copy.copy(self)

  @property
  def label(self) -> str:
    return f"{self.name}_{self.unique_id}"

  def named(self: RASPExprT, name: str) -> RASPExprT:
    """Convenience method for adding a name."""
    return annotate(self, name=name)

  def annotated(self: RASPExprT, **annotations) -> RASPExprT:
    """Convenience method for adding annotations."""
    return annotate(self, **annotations)

def annotate(expr: RASPExprT, **annotations) -> RASPExprT:
  """Creates a new expr with added annotations."""
  new = expr.copy()
  # Note that new annotations will overwrite existing ones with matching keys.
  new.annotations = {**expr.annotations, **annotations}
  return new

### S-Ops.

class SOp(RASPExpr):
  """A Sequence Operation."""

  def __call__(self, xs: Sequence[Value]) -> Sequence[Value]:
    return evaluate(self, xs)  # pytype: disable=bad-return-type

  # Allow construction of SOps using numeric operators with constant values.
  # Note: if inheriting SOp by a dataclass, make sure to disable eq and order,
  # as they will override these.

  def __lt__(self, other: Value) -> "SOp":
    """self < other."""
    return Map(lambda x: x < other, self)

  def __le__(self, other: Value) -> "SOp":
    """self <= other."""
    return Map(lambda x: x <= other, self)

  def __eq__(self, other: Value) -> "SOp":
    """self == other."""
    return Map(lambda x: x == other, self)

  def __ne__(self, other: Value) -> "SOp":
    """self != other."""
    return Map(lambda x: x != other, self)

  def __gt__(self, other: Value) -> "SOp":
    """self > other."""
    return Map(lambda x: x > other, self)

  def __ge__(self, other: Value) -> "SOp":
    """self >= other."""
    return Map(lambda x: x >= other, self)

  def __add__(self, other: Union["SOp", Value]) -> "SOp":
    """self + other."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x + y, self, other)
    return Map(lambda x: x + other, self)

  def __radd__(self, other: Union["SOp", Value]) -> "SOp":
    """other + self."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x + y, other, self)
    return Map(lambda x: other + x, self)

  def __sub__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """self - other."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x - y, self, other)
    return Map(lambda x: x - other, self)

  def __rsub__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """other - self."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x - y, other, self)
    return Map(lambda x: other - x, self)

  def __mul__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """self * other."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x * y, self, other)
    return Map(lambda x: x * other, self)

  def __rmul__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """other * self."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x * y, other, self)
    return Map(lambda x: other * x, self)

  def __truediv__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """self / other."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x / y, self, other)
    return Map(lambda x: x / other, self)

  def __rtruediv__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """other / self."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x / y, other, self)
    return Map(lambda x: other / x, self)

  def __invert__(self) -> "SOp":
    return Map(lambda x: not x, self)

  def __and__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """self & other."""

    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x and y, self, other)
    return Map(lambda x: x and other, self)