Index

B

BaseEmbeddingConfig (class in torchrec.modules.embedding_configs)

C

collective_plan() (torchrec.distributed.planner.planners.EmbeddingShardingPlanner method)

concat() (torchrec.sparse.jagged_tensor.KeyedJaggedTensor static method) copy() (torchrec.distributed.model_parallel.DistributedModelParallel method)

D

device (torchrec.modules.embedding_modules.EmbeddingBagCollection property) (torchrec.modules.embedding_modules.EmbeddingCollection property) device() (torchrec.sparse.jagged_tensor.JaggedTensor method) (torchrec.sparse.jagged_tensor.KeyedJaggedTensor method) (torchrec.sparse.jagged_tensor.KeyedTensor method)

DistributedModelParallel (class in torchrec.distributed.model_parallel)

E

embedding_bag_configs() (torchrec.modules.embedding_modules.EmbeddingBagCollection method) embedding_configs() (torchrec.modules.embedding_modules.EmbeddingCollection method) embedding_dim() (torchrec.modules.embedding_modules.EmbeddingCollection method) embedding_names_by_table() (torchrec.modules.embedding_modules.EmbeddingCollection method) EmbeddingBagCollection (class in torchrec.modules.embedding_modules) EmbeddingBagConfig (class in torchrec.modules.embedding_configs) EmbeddingCollection (class in torchrec.modules.embedding_modules) EmbeddingConfig (class in torchrec.modules.embedding_configs) EmbeddingEnumerator (class in torchrec.distributed.planner.enumerators)

EmbeddingPerfEstimator (class in torchrec.distributed.planner.shard_estimators) EmbeddingShardingPlanner (class in torchrec.distributed.planner.planners) EmbeddingStorageEstimator (class in torchrec.distributed.planner.shard_estimators) empty() (torchrec.sparse.jagged_tensor.JaggedTensor static method) (torchrec.sparse.jagged_tensor.KeyedJaggedTensor static method) empty_like() (torchrec.sparse.jagged_tensor.KeyedJaggedTensor static method) enumerate() (torchrec.distributed.planner.enumerators.EmbeddingEnumerator method) estimate() (torchrec.distributed.planner.shard_estimators.EmbeddingPerfEstimator method) (torchrec.distributed.planner.shard_estimators.EmbeddingStorageEstimator method)

F

feedback() (torchrec.distributed.planner.proposers.GreedyProposer method) forward() (torchrec.distributed.model_parallel.DistributedModelParallel method) (torchrec.modules.embedding_modules.EmbeddingBagCollection method) (torchrec.modules.embedding_modules.EmbeddingCollection method) from_dense() (torchrec.sparse.jagged_tensor.JaggedTensor static method)

from_dense_lengths() (torchrec.sparse.jagged_tensor.JaggedTensor static method) from_jt_dict() (torchrec.sparse.jagged_tensor.KeyedJaggedTensor static method) from_lengths_sync() (torchrec.sparse.jagged_tensor.KeyedJaggedTensor static method) from_offsets_sync() (torchrec.sparse.jagged_tensor.KeyedJaggedTensor static method) from_tensor_list() (torchrec.sparse.jagged_tensor.KeyedTensor static method)

G

get_plan_for_module() (torchrec.distributed.types.ShardingPlan method)

GreedyPerfPartitioner (class in torchrec.distributed.planner.partitioners) GreedyProposer (class in torchrec.distributed.planner.proposers)

H

HeuristicalStorageReservation (class in torchrec.distributed.planner.storage_reservations)

I

index_per_key() (torchrec.sparse.jagged_tensor.KeyedJaggedTensor method) init_data_parallel() (torchrec.distributed.model_parallel.DistributedModelParallel method)

inverse_indices() (torchrec.sparse.jagged_tensor.KeyedJaggedTensor method) inverse_indices_or_none() (torchrec.sparse.jagged_tensor.KeyedJaggedTensor method) is_weighted() (torchrec.modules.embedding_modules.EmbeddingBagCollection method)

J

JaggedTensor (class in torchrec.sparse.jagged_tensor)

K

key_dim() (torchrec.sparse.jagged_tensor.KeyedTensor method) KeyedJaggedTensor (class in torchrec.sparse.jagged_tensor)

KeyedTensor (class in torchrec.sparse.jagged_tensor) keys() (torchrec.sparse.jagged_tensor.KeyedJaggedTensor method) (torchrec.sparse.jagged_tensor.KeyedTensor method)

L

length_per_key() (torchrec.sparse.jagged_tensor.KeyedJaggedTensor method) (torchrec.sparse.jagged_tensor.KeyedTensor method) length_per_key_or_none() (torchrec.sparse.jagged_tensor.KeyedJaggedTensor method) lengths() (torchrec.sparse.jagged_tensor.JaggedTensor method) (torchrec.sparse.jagged_tensor.KeyedJaggedTensor method)

lengths_offset_per_key() (torchrec.sparse.jagged_tensor.KeyedJaggedTensor method) lengths_or_none() (torchrec.sparse.jagged_tensor.JaggedTensor method) (torchrec.sparse.jagged_tensor.KeyedJaggedTensor method) load() (torchrec.distributed.planner.proposers.GreedyProposer method) load_state_dict() (torchrec.distributed.model_parallel.DistributedModelParallel method)

M

module torchrec.distributed.model_parallel torchrec.distributed.planner.enumerators torchrec.distributed.planner.partitioners torchrec.distributed.planner.planners torchrec.distributed.planner.proposers torchrec.distributed.planner.shard_estimators, [1] torchrec.distributed.planner.storage_reservations torchrec.distributed.types torchrec.inference.modules torchrec.modules.embedding_configs torchrec.modules.embedding_modules torchrec.sparse.jagged_tensor

module (torchrec.distributed.model_parallel.DistributedModelParallel property)

N

named_buffers() (torchrec.distributed.model_parallel.DistributedModelParallel method)

named_parameters() (torchrec.distributed.model_parallel.DistributedModelParallel method) need_indices() (torchrec.modules.embedding_modules.EmbeddingCollection method)

O

offset_per_key() (torchrec.sparse.jagged_tensor.KeyedJaggedTensor method) (torchrec.sparse.jagged_tensor.KeyedTensor method) offset_per_key_or_none() (torchrec.sparse.jagged_tensor.KeyedJaggedTensor method)

offsets() (torchrec.sparse.jagged_tensor.JaggedTensor method) (torchrec.sparse.jagged_tensor.KeyedJaggedTensor method) offsets_or_none() (torchrec.sparse.jagged_tensor.JaggedTensor method) (torchrec.sparse.jagged_tensor.KeyedJaggedTensor method)

P

partition() (torchrec.distributed.planner.partitioners.GreedyPerfPartitioner method) perf_func_emb_wall_time() (torchrec.distributed.planner.shard_estimators.EmbeddingPerfEstimator class method) permute() (torchrec.sparse.jagged_tensor.KeyedJaggedTensor method)

plan (torchrec.distributed.types.ShardingPlan attribute) plan() (torchrec.distributed.planner.planners.EmbeddingShardingPlanner method) populate_estimates() (torchrec.distributed.planner.enumerators.EmbeddingEnumerator method) propose() (torchrec.distributed.planner.proposers.GreedyProposer method)

Q

quantize_inference_model() (in module torchrec.inference.modules)

R

record_stream() (torchrec.sparse.jagged_tensor.JaggedTensor method) (torchrec.sparse.jagged_tensor.KeyedJaggedTensor method) (torchrec.sparse.jagged_tensor.KeyedTensor method)

regroup() (torchrec.sparse.jagged_tensor.KeyedTensor static method) regroup_as_dict() (torchrec.sparse.jagged_tensor.KeyedTensor static method) reset_parameters() (torchrec.modules.embedding_modules.EmbeddingBagCollection method) (torchrec.modules.embedding_modules.EmbeddingCollection method)

S

shard_quant_model() (in module torchrec.inference.modules) ShardingPlan (class in torchrec.distributed.types) split() (torchrec.sparse.jagged_tensor.KeyedJaggedTensor method) state_dict() (torchrec.distributed.model_parallel.DistributedModelParallel method)

stride() (torchrec.sparse.jagged_tensor.KeyedJaggedTensor method) stride_per_key() (torchrec.sparse.jagged_tensor.KeyedJaggedTensor method) stride_per_key_per_rank() (torchrec.sparse.jagged_tensor.KeyedJaggedTensor method) sync() (torchrec.sparse.jagged_tensor.KeyedJaggedTensor method)

T

to() (torchrec.sparse.jagged_tensor.JaggedTensor method) (torchrec.sparse.jagged_tensor.KeyedJaggedTensor method) (torchrec.sparse.jagged_tensor.KeyedTensor method) to_dense() (torchrec.sparse.jagged_tensor.JaggedTensor method) to_dense_weights() (torchrec.sparse.jagged_tensor.JaggedTensor method) to_dict() (torchrec.sparse.jagged_tensor.KeyedJaggedTensor method) (torchrec.sparse.jagged_tensor.KeyedTensor method) to_padded_dense() (torchrec.sparse.jagged_tensor.JaggedTensor method) to_padded_dense_weights() (torchrec.sparse.jagged_tensor.JaggedTensor method) torchrec.distributed.model_parallel module torchrec.distributed.planner.enumerators module torchrec.distributed.planner.partitioners module torchrec.distributed.planner.planners module

torchrec.distributed.planner.proposers module torchrec.distributed.planner.shard_estimators module, [1] torchrec.distributed.planner.storage_reservations module torchrec.distributed.types module torchrec.inference.modules module torchrec.modules.embedding_configs module torchrec.modules.embedding_modules module torchrec.sparse.jagged_tensor module

U

unsync() (torchrec.sparse.jagged_tensor.KeyedJaggedTensor method)

V

values() (torchrec.sparse.jagged_tensor.JaggedTensor method) (torchrec.sparse.jagged_tensor.KeyedJaggedTensor method) (torchrec.sparse.jagged_tensor.KeyedTensor method)

variable_stride_per_key() (torchrec.sparse.jagged_tensor.KeyedJaggedTensor method)

W

weights() (torchrec.sparse.jagged_tensor.JaggedTensor method) (torchrec.sparse.jagged_tensor.KeyedJaggedTensor method)

weights_or_none() (torchrec.sparse.jagged_tensor.JaggedTensor method) (torchrec.sparse.jagged_tensor.KeyedJaggedTensor method)