use where instead of mul in Embedding (#13112)

2025-11-05 12:49:01 -05:00
parent edc4e1aede
commit 52f0081e77
1 changed files with 1 additions and 1 deletions
@@ -323,7 +323,7 @@ class Embedding:
    if not dtypes.is_int(idx.dtype): raise TypeError(f"Expected integer dtype for index in embedding, got {idx.dtype}")
    big_shp = idx.shape+(self.vocab_sz, self.embed_sz)
    arange, idx, vals = self.arange.expand(big_shp), idx.reshape(idx.shape+(1, 1)).expand(big_shp), self.weight.expand(big_shp)
-    return (arange == idx).mul(vals).sum(-2, dtype=vals.dtype)
+    return (arange == idx).where(vals, 0).sum(-2, dtype=vals.dtype)

 class LSTMCell:
  """