def make_apply_autoencoder(
    autoencoder_context: AutoencoderContext,
    use_no_grad: bool = True,  # use True to avoid keeping gradient info for autoencoder;
    # TODO: consider deleting in favor of universal non-gradient-keeping at the outside of ScalarDeriver base functions
) -> Callable[[torch.Tensor, LayerIndex], torch.Tensor]:
    """
    Returns a function that takes a tensor of activations and returns a tensor of the autoencoder
    latent representation of each token.
    """
    # TODO(sbills): Resolve the circular import between this file and attention.py.
    from neuron_explainer.activations.derived_scalars.attention import make_reshape_fn

    # reshape activations to be (n_tokens, n_inputs)
    dst = autoencoder_context.dst
    reshape_fn = make_reshape_fn(dst)

    def apply_autoencoder(raw_activations: torch.Tensor, layer_index: LayerIndex) -> torch.Tensor:
        assert (
            layer_index in autoencoder_context.layer_indices
        ), f"Layer index {layer_index} not in {autoencoder_context.layer_indices}"
        autoencoder = autoencoder_context.get_autoencoder(layer_index)
        latent_activations = autoencoder.encode(reshape_fn(raw_activations))
        return latent_activations  # shape (n_tokens, n_latents)

    if use_no_grad:
        apply_autoencoder = torch.no_grad()(apply_autoencoder)

    return apply_autoencoder

    Returns a function that takes a tensor of activations and returns a tensor of the autoencoder
    latent representation of each token.
    
neuron_explainer/activations/derived_scalars/reconstituted.py