Correct typo enbedding -> embedding

jetstream_pt/layers.py

@@ -313,7 +313,7 @@ def create_quantized_from_nn_linear(
   return obj
 
 
-def get_quantized_enbedding_layer(config: "QuantizationConfig"):
+def get_quantized_embedding_layer(config: "QuantizationConfig"):
   if not config.enable_weight_quantization:
     return nn.Embedding
   else:

jetstream_pt/third_party/llama/model_exportable.py

@@ -14,7 +14,7 @@
     RMSNorm,
     WeightOnlyBlockwiseQuantizedLinear,
     WeightOnlyPerChannelQuantizedLinear,
-    get_quantized_enbedding_layer,
+    get_quantized_embedding_layer,
     get_quantized_linear_layer,
 )
 from torch import nn
@@ -188,7 +188,7 @@ def __init__(
     self.vocab_size = params.vocab_size
     self.n_layers = params.n_layers
 
-    Embedding = get_quantized_enbedding_layer(env.quant_config)
+    Embedding = get_quantized_embedding_layer(env.quant_config)
     self.tok_embeddings = Embedding(
         params.vocab_size,
         params.dim,

jetstream_pt/third_party/mixtral/model.py

@@ -21,7 +21,7 @@
 from torch import Tensor
 from torch.nn import functional as F
 from .config import ModelArgs, find_multiple
-from jetstream_pt.layers import Attention, get_quantized_linear_layer, get_quantized_enbedding_layer
+from jetstream_pt.layers import Attention, get_quantized_linear_layer, get_quantized_embedding_layer
 
 import jax
 
@@ -33,7 +33,7 @@ def __init__(self, config: ModelArgs, env) -> None:
     self.config = config
     self.env = env
 
-    Embedding = get_quantized_enbedding_layer(env.quant_config)
+    Embedding = get_quantized_embedding_layer(env.quant_config)
     self.tok_embeddings = Embedding(
         config.vocab_size, config.dim, device=config.device
     )