devinterp package

Subpackages

• devinterp.optim package
  • Submodules
  • devinterp.optim.metrics module
    • Metrics
      • Metrics.DOT_FIELDS
      • Metrics.NORM_FIELDS
      • Metrics.add_dot_products_()
      • Metrics.add_sum_squared_()
      • Metrics.aggregate()
      • Metrics.distance
      • Metrics.dot_grad_noise
      • Metrics.dot_grad_prior
      • Metrics.dot_prior_noise
      • Metrics.localization
      • Metrics.noise
      • Metrics.numel
      • Metrics.prior
      • Metrics.scaled_grad
      • Metrics.sqrt_norms_()
      • Metrics.to()
      • Metrics.unscaled_grad
      • Metrics.weight_decay
      • Metrics.zero_()
  • devinterp.optim.preconditioner module
    • CompositePreconditioner
    • IdentityPreconditioner
    • MaskPreconditioner
    • NHTPreconditioning
    • Preconditioner
      • Preconditioner.get_coefficients()
    • PreconditionerCoefs
      • PreconditionerCoefs.count()
      • PreconditionerCoefs.grad_coef
      • PreconditionerCoefs.grad_correction
      • PreconditionerCoefs.index()
      • PreconditionerCoefs.noise_coef
      • PreconditionerCoefs.overall_coef
      • PreconditionerCoefs.prior_coef
    • RMSpropPreconditioner
  • devinterp.optim.prior module
    • CompositePrior
      • CompositePrior.key
    • GaussianPrior
      • GaussianPrior.grad()
      • GaussianPrior.initialize()
      • GaussianPrior.key
    • Prior
      • Prior.grad()
      • Prior.initialize()
      • Prior.key
    • UniformPrior
      • UniformPrior.key
  • devinterp.optim.sgld module
    • SGLD
      • SGLD.add_param_group()
      • SGLD.get_metrics()
      • SGLD.get_sketches()
      • SGLD.iter_group_metrics()
      • SGLD.load_state_dict()
      • SGLD.register_load_state_dict_post_hook()
      • SGLD.register_load_state_dict_pre_hook()
      • SGLD.register_state_dict_post_hook()
      • SGLD.register_state_dict_pre_hook()
      • SGLD.register_step_post_hook()
      • SGLD.register_step_pre_hook()
      • SGLD.state_dict()
      • SGLD.step()
      • SGLD.zero_grad()
  • devinterp.optim.sgmcmc module
    • SGMCMC
      • SGMCMC.add_param_group()
      • SGMCMC.get_metrics()
      • SGMCMC.get_params()
      • SGMCMC.get_sketches()
      • SGMCMC.iter_group_metrics()
      • SGMCMC.load_state_dict()
      • SGMCMC.register_load_state_dict_post_hook()
      • SGMCMC.register_load_state_dict_pre_hook()
      • SGMCMC.register_state_dict_post_hook()
      • SGMCMC.register_state_dict_pre_hook()
      • SGMCMC.register_step_post_hook()
      • SGMCMC.register_step_pre_hook()
      • SGMCMC.rmsprop_sgld()
      • SGMCMC.sgld()
      • SGMCMC.sgnht()
      • SGMCMC.state_dict()
      • SGMCMC.step()
      • SGMCMC.zero_grad()
  • devinterp.optim.sgnht module
    • SGNHT
      • SGNHT.add_param_group()
      • SGNHT.load_state_dict()
      • SGNHT.register_load_state_dict_post_hook()
      • SGNHT.register_load_state_dict_pre_hook()
      • SGNHT.register_state_dict_post_hook()
      • SGNHT.register_state_dict_pre_hook()
      • SGNHT.register_step_post_hook()
      • SGNHT.register_step_pre_hook()
      • SGNHT.state_dict()
      • SGNHT.zero_grad()
  • devinterp.optim.sketch module
    • CountSketch
      • CountSketch.hash_indices
      • CountSketch.hash_signs
      • CountSketch.scatter_into_()
      • CountSketch.sketch()
    • SketchBuffer
      • SketchBuffer.QUANTITIES
      • SketchBuffer.localization
      • SketchBuffer.noise
      • SketchBuffer.scaled_grad
      • SketchBuffer.unscaled_grad
      • SketchBuffer.weight_decay
  • devinterp.optim.utils module
  • Module contents
• devinterp.slt package
  • Submodules
  • devinterp.slt.bif module
    • bif()
    • compute_bif()
  • devinterp.slt.config module
    • SamplerConfig
      • SamplerConfig.batch_size
      • SamplerConfig.bounding_box_size
      • SamplerConfig.copy()
      • SamplerConfig.epoch_mode
      • SamplerConfig.gradient_accumulation_steps
      • SamplerConfig.init_noise
      • SamplerConfig.init_seed
      • SamplerConfig.llc_weight_decay
      • SamplerConfig.localization
      • SamplerConfig.lr
      • SamplerConfig.match_sampling_input_ids_across_chains
      • SamplerConfig.model_computed_fields
      • SamplerConfig.model_config
      • SamplerConfig.model_construct()
      • SamplerConfig.model_copy()
      • SamplerConfig.model_dump()
      • SamplerConfig.model_dump_json()
      • SamplerConfig.model_extra
      • SamplerConfig.model_fields
      • SamplerConfig.model_fields_set
      • SamplerConfig.model_json_schema()
      • SamplerConfig.model_parametrized_name()
      • SamplerConfig.model_post_init()
      • SamplerConfig.model_rebuild()
      • SamplerConfig.model_validate()
      • SamplerConfig.model_validate_json()
      • SamplerConfig.model_validate_strings()
      • SamplerConfig.n_beta
      • SamplerConfig.noise_level
      • SamplerConfig.num_burnin_steps
      • SamplerConfig.num_chains
      • SamplerConfig.num_draws
      • SamplerConfig.num_init_loss_batches
      • SamplerConfig.num_steps_bw_draws
      • SamplerConfig.sampling_method
      • SamplerConfig.sampling_method_kwargs
      • SamplerConfig.save_metrics
      • SamplerConfig.shuffle
  • devinterp.slt.covariance module
    • batch_corrcoef()
    • batch_cov()
    • xr_corrcoef_with_torch_backend()
  • devinterp.slt.llc module
    • compute_llc()
    • llc()
  • devinterp.slt.lm_loss module
    • NonFiniteLogitsError
      • NonFiniteLogitsError.add_note()
      • NonFiniteLogitsError.with_traceback()
    • compute_per_token_loss()
    • lm_cross_entropy_loss()
    • lm_forward_logits()
    • make_evaluate_fn()
  • devinterp.slt.observables module
    • DeterministicShuffledSampler
    • Observable
      • Observable.obs_id
      • Observable.input_ids
      • Observable.n_samples
      • Observable.context_length
      • Observable.compute_loss()
    • to_obs_id()
  • devinterp.slt.sampler module
    • MicroCallback
    • sample_single_chain()
  • devinterp.slt.sampling module
    • sample()
  • devinterp.slt.susceptibilities module
    • compute_susceptibilities()
    • susceptibilities()
  • devinterp.slt.weight_restrictions module
    • create_param_masks()
    • preview_weight_restriction()
  • devinterp.slt.writing module
    • ZarrWriter
      • ZarrWriter.arrays
      • ZarrWriter.buffer_device
      • ZarrWriter.chain_buffer_size
      • ZarrWriter.push()
      • ZarrWriter.write()
  • devinterp.slt.zarr_schema module
    • DataArraySpec
      • DataArraySpec.dims
      • DataArraySpec.dtype_str
      • DataArraySpec.shape
    • ZarrSchema
      • ZarrSchema.arrays_meta
      • ZarrSchema.create_hierarchy()
      • ZarrSchema.group_attrs
  • Module contents

Submodules

devinterp.utils module

devinterp.utils.tokenize_and_concatenate(dataset, tokenizer, streaming: bool = False, max_length: int = 1024, column_name: str = 'text', add_bos_token: bool = True, num_proc: int = 10)

Tokenize and concatenate a text dataset into fixed-length sequences.

Based on TransformerLens (MIT License, Copyright 2022 TransformerLensOrg): https://github.com/TransformerLensOrg/TransformerLens Core algorithm unchanged, with local additions: input validation (bos token and max_length checks), numpy reshape in place of einops, and the output column renamed from “tokens” to “input_ids”.

Joins all text separated by EOS tokens, tokenizes, then reshapes into (num_sequences, max_length) chunks.

Parameters:

• dataset – HuggingFace text dataset.

• tokenizer – HuggingFace tokenizer with bos_token_id and eos_token_id.

• streaming – If True, disables parallel tokenization.

• max_length – Context window length.

• column_name – Name of the text column.

• add_bos_token – Prepend BOS to each sequence (reduces usable length by 1).

• num_proc – Number of processes for dataset.map().

Returns:

Dataset with an “input_ids” column of torch tensors.

Module contents