|
|
| from __future__ import annotations |
| import hashlib |
| import math |
| import os |
| from collections.abc import Callable, Mapping |
| from dataclasses import dataclass |
| from numbers import Integral, Real |
| from pathlib import Path |
| import torch |
| import torch.nn.functional as F |
| PINNED_ARTIFACT_VERSION = 1 |
| PINNED_ARTIFACT_ALGORITHM = 'recursive_centered_pca_rank_median_v1' |
| PINNED_ARTIFACT_METRIC = 'normalized_squared_l2_tree_ultrametric' |
| PINNED_PRODUCTION_VOCAB_SIZE = 131072 |
| PINNED_PRODUCTION_DEPTH = 17 |
| DEFAULT_SEMANTIC_TREE_STEPS = 4 |
| DEFAULT_SEMANTIC_TREE_SCHEDULE = 'shift' |
| DEFAULT_SEMANTIC_TREE_SHIFT = 5.0 |
| MAX_SEMANTIC_TREE_REFINEMENT_GAP = 2 |
| SUPPORTED_SEMANTIC_TREE_INFERENCE_SCHEDULES = frozenset({'linear', 'shift', 'logit_normal', 'cosine', 'cosine2', 'exp'}) |
|
|
| def _plain_positive_int(value: object, *, name: str) -> int: |
| if isinstance(value, bool) or not isinstance(value, Integral): |
| raise TypeError(f'{name} must be a positive integer') |
| value = int(value) |
| if value < 1: |
| raise ValueError(f'{name} must be a positive integer') |
| return value |
|
|
| def _plain_nonnegative_int(value: object, *, name: str) -> int: |
| if isinstance(value, bool) or not isinstance(value, Integral): |
| raise TypeError(f'{name} must be a non-negative integer') |
| value = int(value) |
| if value < 0: |
| raise ValueError(f'{name} must be a non-negative integer') |
| return value |
|
|
| def _plain_positive_float(value: object, *, name: str) -> float: |
| if isinstance(value, bool) or not isinstance(value, Real): |
| raise TypeError(f'{name} must be a finite positive number') |
| value = float(value) |
| if not math.isfinite(value) or value <= 0.0: |
| raise ValueError(f'{name} must be a finite positive number') |
| return value |
|
|
| def _is_power_of_two(value: int) -> bool: |
| return value > 0 and value & value - 1 == 0 |
|
|
| def resolve_semantic_tree_steps(config, override=None): |
| steps = 4 if override is None else override |
| if isinstance(steps, bool) or steps not in (1, 4, 8, 16): |
| raise ValueError('refinement_steps must be 1, 4, 8, or 16') |
| return int(steps) |
|
|
| @dataclass(frozen=True, slots=True) |
| class SemanticTreeArtifactContract: |
| path: str | os.PathLike[str] |
| bytes: int | None = None |
| sha256: str | None = None |
| algorithm: str | None = PINNED_ARTIFACT_ALGORITHM |
|
|
| def __post_init__(self) -> None: |
| if not isinstance(self.path, (str, os.PathLike)): |
| raise TypeError('semantic-tree artifact path must be path-like') |
| raw_path = os.fspath(self.path) |
| if not isinstance(raw_path, str) or not raw_path.strip(): |
| raise ValueError('semantic-tree artifact path must be nonempty') |
| object.__setattr__(self, 'path', Path(raw_path).expanduser()) |
| if self.bytes is not None: |
| object.__setattr__(self, 'bytes', _plain_positive_int(self.bytes, name='semantic-tree artifact bytes')) |
| if self.sha256 is not None: |
| if not isinstance(self.sha256, str): |
| raise TypeError('semantic-tree artifact sha256 must be a string') |
| digest = self.sha256.strip().lower() |
| if len(digest) != 64 or any((character not in '0123456789abcdef' for character in digest)): |
| raise ValueError('semantic-tree artifact sha256 must be 64 hexadecimal characters') |
| object.__setattr__(self, 'sha256', digest) |
| if self.algorithm not in (None, PINNED_ARTIFACT_ALGORITHM): |
| raise ValueError(f'Unsupported semantic-tree algorithm: {self.algorithm!r}') |
|
|
| def production_semantic_tree_contract(path: str | os.PathLike[str]) -> SemanticTreeArtifactContract: |
| return SemanticTreeArtifactContract(path=path, sha256='4b7ee94fbb10fd4146a4ad167b595b037e3f591eaa12ecf3f0b411c82384a714', algorithm=PINNED_ARTIFACT_ALGORITHM) |
|
|
| def _sha256_file(path: Path) -> str: |
| digest = hashlib.sha256() |
| with path.open('rb') as artifact_file: |
| for block in iter(lambda : artifact_file.read(1024 * 1024), b''): |
| digest.update(block) |
| return digest.hexdigest() |
|
|
| class BalancedSemanticTree: |
|
|
| def __init__(self, leaf_to_token: torch.Tensor) -> None: |
| if not torch.is_tensor(leaf_to_token): |
| raise TypeError('leaf_to_token must be a tensor') |
| if leaf_to_token.ndim != 1: |
| raise ValueError('leaf_to_token must be one-dimensional') |
| if leaf_to_token.dtype == torch.bool or leaf_to_token.is_floating_point() or leaf_to_token.is_complex(): |
| raise TypeError('leaf_to_token must contain integer token IDs') |
| vocab_size = int(leaf_to_token.numel()) |
| if not _is_power_of_two(vocab_size): |
| raise ValueError('semantic-tree vocabulary size must be a power of two') |
| permutation = leaf_to_token.to(dtype=torch.long).contiguous() |
| expected = torch.arange(vocab_size, device=permutation.device) |
| if not torch.equal(torch.sort(permutation).values, expected): |
| raise ValueError('leaf_to_token must be a permutation of [0, V)') |
| token_to_leaf = torch.empty_like(permutation) |
| token_to_leaf.scatter_(0, permutation, expected) |
| self.vocab_size = vocab_size |
| self.depth = int(math.log2(vocab_size)) |
| self.leaf_to_token = permutation |
| self.token_to_leaf = token_to_leaf |
|
|
| @classmethod |
| def _from_validated(cls, leaf_to_token: torch.Tensor, token_to_leaf: torch.Tensor) -> BalancedSemanticTree: |
| tree = cls.__new__(cls) |
| tree.vocab_size = int(leaf_to_token.numel()) |
| tree.depth = int(math.log2(tree.vocab_size)) |
| tree.leaf_to_token = leaf_to_token |
| tree.token_to_leaf = token_to_leaf |
| return tree |
|
|
| @classmethod |
| def from_artifact(cls, contract: SemanticTreeArtifactContract, expected_vocab_size: int=PINNED_PRODUCTION_VOCAB_SIZE) -> BalancedSemanticTree: |
| if not isinstance(contract, SemanticTreeArtifactContract): |
| raise TypeError('contract must be a SemanticTreeArtifactContract') |
| expected_vocab_size = _plain_positive_int(expected_vocab_size, name='expected_vocab_size') |
| if not _is_power_of_two(expected_vocab_size): |
| raise ValueError('expected_vocab_size must be a power of two') |
| path = contract.path |
| assert isinstance(path, Path) |
| if not path.is_file(): |
| raise FileNotFoundError(f'semantic-tree artifact is not a file: {path}') |
| actual_bytes = path.stat().st_size |
| if contract.bytes is not None and actual_bytes != contract.bytes: |
| raise ValueError(f'semantic-tree artifact byte count mismatch: expected {contract.bytes}, got {actual_bytes}') |
| if contract.sha256 is not None: |
| actual_sha256 = _sha256_file(path) |
| if actual_sha256 != contract.sha256: |
| raise ValueError(f'semantic-tree artifact sha256 mismatch: expected {contract.sha256}, got {actual_sha256}') |
| artifact = torch.load(path, map_location='cpu', weights_only=True) |
| if not isinstance(artifact, Mapping): |
| raise TypeError('semantic-tree artifact must contain a mapping') |
| algorithm = artifact.get('algorithm') |
| allowed_algorithms = (contract.algorithm,) if contract.algorithm is not None else (PINNED_ARTIFACT_ALGORITHM,) |
| if algorithm not in allowed_algorithms: |
| raise ValueError(f'semantic-tree artifact algorithm mismatch: expected one of {allowed_algorithms!r}, got {algorithm!r}') |
| pinned_values = {'version': PINNED_ARTIFACT_VERSION, 'metric': PINNED_ARTIFACT_METRIC, 'vocab_size': expected_vocab_size, 'depth': int(math.log2(expected_vocab_size)), 'radii_monotone': True} |
| for (field_name, expected_value) in pinned_values.items(): |
| if artifact.get(field_name) != expected_value: |
| raise ValueError(f'semantic-tree artifact {field_name} mismatch: expected {expected_value!r}, got {artifact.get(field_name)!r}') |
| leaf_to_token = artifact.get('leaf_to_token') |
| heap_radii = artifact.get('heap_radii') |
| if not torch.is_tensor(leaf_to_token) or leaf_to_token.dtype != torch.int32: |
| raise TypeError('semantic-tree leaf_to_token must be torch.int32') |
| if tuple(leaf_to_token.shape) != (expected_vocab_size,): |
| raise ValueError('semantic-tree leaf_to_token shape disagrees with vocabulary') |
| if not torch.is_tensor(heap_radii) or heap_radii.dtype != torch.float32: |
| raise TypeError('semantic-tree heap_radii must be torch.float32') |
| if tuple(heap_radii.shape) != (expected_vocab_size - 1,): |
| raise ValueError('semantic-tree heap_radii shape disagrees with vocabulary') |
| if not bool(torch.isfinite(heap_radii).all()) or bool((heap_radii < 0).any()): |
| raise ValueError('semantic-tree heap_radii must be finite and nonnegative') |
| return cls(leaf_to_token) |
|
|
| @property |
| def device(self) -> torch.device: |
| return self.leaf_to_token.device |
|
|
| def to(self, device: torch.device | str) -> BalancedSemanticTree: |
| resolved = torch.device(device) |
| if resolved == self.device: |
| return self |
| return type(self)._from_validated(self.leaf_to_token.to(resolved), self.token_to_leaf.to(resolved)) |
|
|
| def _check_tokens(self, tokens: torch.Tensor, *, name: str) -> None: |
| if not torch.is_tensor(tokens): |
| raise TypeError(f'{name} must be a tensor') |
| if tokens.dtype == torch.bool or tokens.is_floating_point() or tokens.is_complex(): |
| raise TypeError(f'{name} must contain integer code IDs') |
| if tokens.device != self.device: |
| raise ValueError(f'{name} and semantic-tree lookups must share a device') |
| if tokens.numel() and bool(((tokens < 0) | (tokens >= self.vocab_size)).any()): |
| raise ValueError(f'{name} contains code IDs outside [0, V)') |
|
|
| def _levels(self, levels: int | torch.Tensor, shape: torch.Size, device: torch.device) -> torch.Tensor: |
| if isinstance(levels, bool) or (not isinstance(levels, Integral) and (not torch.is_tensor(levels))): |
| raise TypeError('levels must contain integers') |
| level_tensor = torch.as_tensor(levels, device=device) |
| if level_tensor.dtype == torch.bool or level_tensor.is_floating_point() or level_tensor.is_complex(): |
| raise TypeError('levels must contain integers') |
| if level_tensor.ndim < len(shape): |
| level_tensor = level_tensor.reshape(level_tensor.shape + (1,) * (len(shape) - level_tensor.ndim)) |
| try: |
| level_tensor = torch.broadcast_to(level_tensor, shape).long() |
| except RuntimeError as error: |
| raise ValueError('levels are not broadcastable to endpoint shape') from error |
| if bool((level_tensor < 0).any()) or bool((level_tensor > self.depth).any()): |
| raise ValueError(f'levels must lie in [0, {self.depth}]') |
| return level_tensor |
|
|
| def node_ids(self, tokens: torch.Tensor, levels: int | torch.Tensor) -> torch.Tensor: |
| self._check_tokens(tokens, name='tokens') |
| level_tensor = self._levels(levels, tokens.shape, tokens.device) |
| leaves = self.token_to_leaf[tokens.long()] |
| return torch.bitwise_right_shift(leaves, self.depth - level_tensor) |
|
|
| @torch.no_grad() |
| def sample_representatives(self, endpoints: torch.Tensor, levels: int | torch.Tensor, generator: torch.Generator | None=None) -> torch.Tensor: |
| self._check_tokens(endpoints, name='endpoints') |
| level_tensor = self._levels(levels, endpoints.shape, endpoints.device) |
| leaves = self.token_to_leaf[endpoints.long()] |
| noisy = level_tensor < self.depth |
| if not bool(noisy.any()): |
| return endpoints.clone() |
| noisy_leaves = leaves[noisy] |
| noisy_levels = level_tensor[noisy] |
| shifts = self.depth - noisy_levels |
| subtree_starts = torch.bitwise_left_shift(torch.bitwise_right_shift(noisy_leaves, shifts), shifts) |
| subtree_sizes = torch.bitwise_left_shift(torch.ones_like(shifts), shifts) |
| random_bits = torch.randint(self.vocab_size, (int(noisy.sum().item()),), device=endpoints.device, dtype=torch.long, generator=generator) |
| sampled_leaves = subtree_starts + torch.bitwise_and(random_bits, subtree_sizes - 1) |
| result = endpoints.clone() |
| result[noisy] = self.leaf_to_token[sampled_leaves].to(endpoints.dtype) |
| return result |
|
|
| def _validate_ages(ages: torch.Tensor, *, steps: int) -> tuple[torch.Tensor, int]: |
| steps = _plain_positive_int(steps, name='steps') |
| if not torch.is_tensor(ages): |
| raise TypeError('ages must be a tensor') |
| if ages.dtype == torch.bool or ages.is_floating_point() or ages.is_complex(): |
| raise TypeError('ages must contain integers') |
| if bool((ages < 0).any()) or bool((ages > steps).any()): |
| raise ValueError(f'ages must lie in [0, {steps}]') |
| return (ages.long(), steps) |
|
|
| def semantic_tree_levels_for_ages(ages: torch.Tensor, steps: int=DEFAULT_SEMANTIC_TREE_STEPS, depth: int=PINNED_PRODUCTION_DEPTH) -> torch.Tensor: |
| (ages, steps) = _validate_ages(ages, steps=steps) |
| depth = _plain_positive_int(depth, name='depth') |
| if steps > depth: |
| raise ValueError('steps cannot exceed tree depth') |
| return torch.div(ages * depth, steps, rounding_mode='floor') |
|
|
| def semantic_tree_noise_for_ages(ages: torch.Tensor, steps: int=DEFAULT_SEMANTIC_TREE_STEPS) -> torch.Tensor: |
| (ages, steps) = _validate_ages(ages, steps=steps) |
| return 1.0 - ages.to(dtype=torch.float32) / float(steps) |
|
|
| def semantic_tree_activation_counts(content_tokens: int, *, activation_steps: int, schedule: str=DEFAULT_SEMANTIC_TREE_SCHEDULE, shift: float=DEFAULT_SEMANTIC_TREE_SHIFT, device: torch.device | str | None=None) -> torch.Tensor: |
| content_tokens = _plain_positive_int(content_tokens, name='content_tokens') |
| activation_steps = min(_plain_positive_int(activation_steps, name='activation_steps'), content_tokens) |
| if not isinstance(schedule, str): |
| raise TypeError('schedule must be a string') |
| schedule = schedule.strip().lower() |
| if schedule not in SUPPORTED_SEMANTIC_TREE_INFERENCE_SCHEDULES: |
| allowed = ', '.join(sorted(SUPPORTED_SEMANTIC_TREE_INFERENCE_SCHEDULES)) |
| raise ValueError(f'unsupported semantic-tree inference schedule {schedule!r}; expected one of: {allowed}') |
| shift = _plain_positive_float(shift, name='shift') |
| sample_times = torch.linspace(0.0, 1.0, activation_steps + 1, dtype=torch.float32) |
| if schedule == 'linear': |
| sigmas = sample_times |
| elif schedule == 'shift': |
| sigmas = shift * sample_times / (1.0 + (shift - 1.0) * sample_times) |
| elif schedule == 'logit_normal': |
| sigmas = 0.5 * (1.0 + torch.erf(torch.logit(sample_times) / math.sqrt(2.0))) |
| elif schedule == 'cosine': |
| sigmas = 1.0 - 0.5 * (1.0 + torch.cos(math.pi * sample_times)) |
| elif schedule == 'cosine2': |
| sigmas = torch.sin(0.5 * math.pi * sample_times) |
| else: |
| sigmas = (1.0 - torch.exp(-5.0 * sample_times)) / (1.0 - math.exp(-5.0)) |
| sigmas = sigmas.clamp(0.0001, 1.0) |
| boundaries = (sigmas * content_tokens).to(dtype=torch.long) |
| counts = (boundaries[1:] - boundaries[:-1]).clamp_min(1) |
| excess = int(counts.sum().item()) - content_tokens |
| cursor = 0 |
| while excess > 0: |
| cursor %= counts.numel() |
| if int(counts[cursor].item()) == 1: |
| cursor += 1 |
| continue |
| counts[cursor] -= 1 |
| excess -= 1 |
| cursor += 1 |
| if int(counts.sum().item()) != content_tokens: |
| raise RuntimeError('semantic-tree activation counts do not cover content') |
| return counts.flip(0).to(device=device) |
|
|
| def resolve_semantic_tree_atomic_block_size(value: object=1) -> int: |
| block_size = _plain_positive_int(value, name='semantic_tree_atomic_block_size') |
| if block_size not in {1, 2, 4, 8}: |
| raise ValueError('semantic_tree_atomic_block_size currently supports 1, 2, 4, or 8') |
| return block_size |
|
|
| def resolve_semantic_tree_refinement_gap(value: object=0) -> int: |
| gap = _plain_nonnegative_int(value, name='semantic_tree_refinement_gap') |
| if gap > MAX_SEMANTIC_TREE_REFINEMENT_GAP: |
| raise ValueError('semantic_tree_refinement_gap currently supports only 0, 1, or 2') |
| return gap |
|
|
| def semantic_tree_decode_schedule(activation_counts: torch.Tensor, *, refinement_steps: int, refinement_gap: int=0) -> torch.Tensor: |
| if not torch.is_tensor(activation_counts): |
| raise TypeError('activation_counts must be a tensor') |
| if activation_counts.ndim != 2: |
| raise ValueError('activation_counts must have shape [batch, rounds]') |
| if activation_counts.dtype not in {torch.uint8, torch.int8, torch.int16, torch.int32, torch.int64}: |
| raise TypeError('activation_counts must use an integer dtype') |
| if activation_counts.shape[1] == 0: |
| raise ValueError('activation_counts must contain at least one round') |
| if bool((activation_counts < 0).any()): |
| raise ValueError('activation_counts must be non-negative') |
| refinement_steps = _plain_positive_int(refinement_steps, name='refinement_steps') |
| refinement_gap = resolve_semantic_tree_refinement_gap(refinement_gap) |
| flush_rounds = refinement_steps - 1 + refinement_steps * refinement_gap |
| if flush_rounds == 0: |
| return activation_counts |
| flush = torch.zeros(activation_counts.shape[0], flush_rounds, dtype=activation_counts.dtype, device=activation_counts.device) |
| return torch.cat((activation_counts, flush), dim=1) |
|
|
| def semantic_tree_refinement_due(ages: torch.Tensor, next_refinement_round: torch.Tensor, *, decode_round: int, refinement_steps: int) -> torch.Tensor: |
| if not torch.is_tensor(ages) or not torch.is_tensor(next_refinement_round): |
| raise TypeError('ages and next_refinement_round must be tensors') |
| if ages.shape != next_refinement_round.shape: |
| raise ValueError('ages and next_refinement_round must align') |
| if ages.device != next_refinement_round.device: |
| raise ValueError('ages and next_refinement_round must share a device') |
| for (tensor, name) in ((ages, 'ages'), (next_refinement_round, 'next_refinement_round')): |
| if tensor.dtype == torch.bool or tensor.is_floating_point(): |
| raise TypeError(f'{name} must use an integer dtype') |
| decode_round = _plain_nonnegative_int(decode_round, name='decode_round') |
| refinement_steps = _plain_positive_int(refinement_steps, name='refinement_steps') |
| return ages.ge(0) & ages.lt(refinement_steps) & next_refinement_round.le(decode_round) |
| EmbeddingSource = torch.Tensor | Callable[[torch.Tensor], torch.Tensor] |
|
|
| def _embed_ids(source: EmbeddingSource, token_ids: torch.Tensor) -> torch.Tensor: |
| if torch.is_tensor(source): |
| if source.ndim != 2 or not source.is_floating_point(): |
| raise TypeError('embedding weight must be a floating-point matrix') |
| if source.device != token_ids.device: |
| raise ValueError('embedding weight and token IDs must share a device') |
| if token_ids.numel() and bool(((token_ids < 0) | (token_ids >= source.shape[0])).any()): |
| raise ValueError('token ID lies outside the embedding weight') |
| return F.embedding(token_ids.long(), source) |
| if not callable(source): |
| raise TypeError('embedding source must be a weight tensor or callable') |
| embedded = source(token_ids) |
| if not torch.is_tensor(embedded): |
| raise TypeError('embedding callback must return a tensor') |
| return embedded |
|
|
| def build_semantic_tree_embeddings(representatives: torch.Tensor, noise_levels: float | torch.Tensor, embedding_callback: EmbeddingSource, mask_token_id: int) -> torch.Tensor: |
| if not torch.is_tensor(representatives): |
| raise TypeError('representatives must be a tensor') |
| if representatives.dtype == torch.bool or representatives.is_floating_point() or representatives.is_complex(): |
| raise TypeError('representatives must contain integer code IDs') |
| if isinstance(mask_token_id, bool) or not isinstance(mask_token_id, Integral): |
| raise TypeError('mask_token_id must be an integer') |
| mask_token_id = int(mask_token_id) |
| if mask_token_id < 0: |
| raise ValueError('mask_token_id must be nonnegative') |
| representative_embeddings = _embed_ids(embedding_callback, representatives) |
| expected_prefix = tuple(representatives.shape) |
| if representative_embeddings.ndim != representatives.ndim + 1 or tuple(representative_embeddings.shape[:-1]) != expected_prefix or (not representative_embeddings.is_floating_point()): |
| raise ValueError('representative embeddings must have shape representatives.shape + [H]') |
| mask_ids = torch.tensor([mask_token_id], dtype=representatives.dtype, device=representatives.device) |
| mask_embedding = _embed_ids(embedding_callback, mask_ids) |
| if mask_embedding.shape != (1, representative_embeddings.shape[-1]): |
| raise ValueError('MASK embedding must have shape [1, H]') |
| if mask_embedding.device != representative_embeddings.device or mask_embedding.dtype != representative_embeddings.dtype: |
| raise ValueError('representative and MASK embeddings must share dtype/device') |
| noise = torch.as_tensor(noise_levels, device=representatives.device, dtype=torch.float32) |
| if noise.ndim < representatives.ndim: |
| noise = noise.reshape(noise.shape + (1,) * (representatives.ndim - noise.ndim)) |
| try: |
| noise = torch.broadcast_to(noise, representatives.shape) |
| except RuntimeError as error: |
| raise ValueError('noise_levels are not broadcastable to representatives') from error |
| if not bool(torch.isfinite(noise).all()) or bool((noise < 0).any()) or bool((noise > 1).any()): |
| raise ValueError('noise_levels must be finite and lie in [0, 1]') |
| noise_in_embedding_dtype = noise.to(dtype=representative_embeddings.dtype) |
| combined = representative_embeddings + noise_in_embedding_dtype.unsqueeze(-1) * mask_embedding |
| return torch.where(noise.eq(0).unsqueeze(-1), representative_embeddings, combined) |
|
|
| @dataclass(frozen=True, slots=True) |
| class SemanticTreeAdvance: |
| representatives: torch.Tensor |
| ages: torch.Tensor |
| levels: torch.Tensor |
| noise_levels: torch.Tensor |
| committed: torch.Tensor |
|
|
| @torch.no_grad() |
| def advance_revisable_semantic_tree(current_representatives: torch.Tensor, x0_proposals: torch.Tensor, ages: torch.Tensor, tree: BalancedSemanticTree, steps: int=DEFAULT_SEMANTIC_TREE_STEPS, generator: torch.Generator | None=None) -> SemanticTreeAdvance: |
| if not isinstance(tree, BalancedSemanticTree): |
| raise TypeError('tree must be a BalancedSemanticTree') |
| tree._check_tokens(current_representatives, name='current_representatives') |
| tree._check_tokens(x0_proposals, name='x0_proposals') |
| if current_representatives.shape != x0_proposals.shape: |
| raise ValueError('current_representatives and x0_proposals must align') |
| if current_representatives.dtype != x0_proposals.dtype: |
| raise ValueError('current_representatives and x0_proposals must share dtype') |
| if not torch.is_tensor(ages) or ages.shape != current_representatives.shape: |
| raise ValueError('ages must align with representative tokens') |
| (validated_ages, steps) = _validate_ages(ages, steps=steps) |
| if validated_ages.device != current_representatives.device: |
| raise ValueError('ages and representative tokens must share a device') |
| if steps > tree.depth: |
| raise ValueError('steps cannot exceed tree depth') |
| unfinished = validated_ages < steps |
| next_ages = torch.where(unfinished, validated_ages + 1, validated_ages) |
| next_levels = semantic_tree_levels_for_ages(next_ages, steps=steps, depth=tree.depth) |
| proposal_representatives = tree.sample_representatives(x0_proposals, next_levels, generator=generator) |
| next_representatives = torch.where(unfinished, proposal_representatives, current_representatives) |
| return SemanticTreeAdvance(representatives=next_representatives, ages=next_ages.to(dtype=ages.dtype), levels=next_levels, noise_levels=semantic_tree_noise_for_ages(next_ages, steps=steps), committed=next_ages.eq(steps)) |
|
|
