ai-edge-torch-nightly 0.7.0.dev20251017__py3-none-any.whl → 0.7.0.dev20251018__py3-none-any.whl

ai_edge_torch/odml_torch/lowerings/_basic.py

@@ -51,6 +51,75 @@ def _aten_mul_tensor(lctx, self: ir.Value, other: ir.Value):
   return stablehlo.multiply(self, other)
 
 
+def _hann_window_impl(
+    lctx: LoweringContext,
+    size: int,
+    periodic: bool,
+    dtype: Optional[torch.dtype],
+) -> ir.Value:
+  if dtype is None:
+    ir_dtype = ir.F32Type.get()
+  else:
+    ir_dtype = utils.torch_dtype_to_ir_element_type(dtype)
+
+  if not isinstance(ir_dtype, ir.FloatType):
+    raise ValueError("hann_window only supports float dtypes.")
+
+  if size == 0:
+    return stablehlo.ConstantOp(
+        ir.RankedTensorType.get((0,), ir_dtype),
+        ir.DenseElementsAttr.get_empty(ir.RankedTensorType.get((0,), ir_dtype)),
+    ).result
+  if size == 1:
+    return utils.splat(1.0, ir_dtype, [1])
+
+  denom = size if periodic else size - 1
+
+  i64 = ir.IntegerType.get_signless(64)
+  iota_type = ir.RankedTensorType.get((size,), i64)
+  n_i64 = stablehlo.IotaOp(
+      iota_type, iota_dimension=ir.IntegerAttr.get(i64, 0)
+  ).result
+
+  n_type = ir.RankedTensorType.get((size,), ir_dtype)
+  n = stablehlo.convert(n_type, n_i64)
+
+  pi_val = math.pi
+  scale = 2.0 * pi_val / denom
+
+  scale_splat = utils.splat(scale, ir_dtype, [size])
+  arg_cos = stablehlo.multiply(n, scale_splat)
+  cos_val = stablehlo.cosine(arg_cos)
+
+  half_splat = utils.splat(0.5, ir_dtype, [size])
+  scaled_cos = stablehlo.multiply(half_splat, cos_val)
+  return stablehlo.subtract(half_splat, scaled_cos)
+
+
+# hann_window(int size, *, ScalarType? dtype=None) -> Tensor
+@lower(torch.ops.aten.hann_window.default)
+def _aten_hann_window_default(
+    lctx: LoweringContext,
+    size: int,
+    *,
+    dtype: Optional[torch.dtype] = None,
+) -> ir.Value:
+  return _hann_window_impl(lctx, size, True, dtype)
+
+
+# hann_window.periodic(int size, bool periodic, *, ScalarType? dtype=None) ->
+# Tensor
+@lower(torch.ops.aten.hann_window.periodic)
+def _aten_hann_window_periodic(
+    lctx: LoweringContext,
+    size: int,
+    periodic: bool,
+    *,
+    dtype: Optional[torch.dtype] = None,
+) -> ir.Value:
+  return _hann_window_impl(lctx, size, periodic, dtype)
+
+
 # cat(Tensor[] tensors, int dim=0) -> Tensor
 # @lower(torch.ops.aten.cat)
 def _aten_cat(lctx, tensors: list[ir.Value], dim: int = 1):

ai_edge_torch/odml_torch/lowerings/_jax_lowerings.py

@@ -541,3 +541,46 @@ def _aten_topk(
     return values, indices
 
   return jax_lowering(lctx, self, k)
+
+
+@registry.lower(torch.ops.aten.multinomial)
+def _aten_multinomial(
+    lctx: LoweringContext,
+    self,
+    num_samples,
+    replacement=False,
+    generator=None,
+):
+  _log_usage(torch.ops.aten.multinomial)
+
+  @jax_bridge.wrap
+  def jax_lowering(self, num_samples, replacement):
+    if generator is not None:
+      logging.warning("aten.multinomial lowering ignores `generator`.")
+
+    assert (
+        num_samples <= self.shape[-1] or replacement
+    ), "cannot take a larger sample than population when replacement=False"
+
+    # TODO: Add proper PRNG key handling.
+    key = jax.random.PRNGKey(0)
+    if self.ndim == 1:
+      return jax.random.choice(
+          key, self.shape[-1], (num_samples,), replace=replacement, p=self
+      ).astype(jnp.int64)
+    else:
+      return jnp.array(
+          [
+              jax.random.choice(
+                  key,
+                  self.shape[-1],
+                  (num_samples,),
+                  replace=replacement,
+                  p=self[i, :],
+              )
+              for i in range(self.shape[0])
+          ],
+          dtype=jnp.int64,
+      )
+
+  return jax_lowering(lctx, self, num_samples, replacement)

ai_edge_torch/version.py

@@ -15,4 +15,4 @@
 
 # The next version of ai-edge-torch.
 # The minor version code should be bumped after every release.
-__version__ = "0.7.0.dev20251017"
+__version__ = "0.7.0.dev20251018"

{ai_edge_torch_nightly-0.7.0.dev20251017.dist-info → ai_edge_torch_nightly-0.7.0.dev20251018.dist-info}/METADATA

@@ -1,6 +1,6 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.4
 Name: ai-edge-torch-nightly
-Version: 0.7.0.dev20251017
+Version: 0.7.0.dev20251018
 Summary: Supporting PyTorch models with the Google AI Edge TFLite runtime.
 Home-page: https://github.com/google-ai-edge/ai-edge-torch
 Keywords: On-Device ML,AI,Google,TFLite,PyTorch,LLMs,GenAI
@@ -37,7 +37,17 @@ Requires-Dist: ai-edge-quantizer-nightly
 Requires-Dist: jax
 Requires-Dist: torch-xla2[odml]>=0.0.1.dev20241201
 Provides-Extra: torch-xla
-Requires-Dist: torch-xla>=2.4.0; extra == "torch-xla"
+Requires-Dist: torch_xla>=2.4.0; extra == "torch-xla"
+Dynamic: classifier
+Dynamic: description
+Dynamic: description-content-type
+Dynamic: home-page
+Dynamic: keywords
+Dynamic: license-file
+Dynamic: provides-extra
+Dynamic: requires-dist
+Dynamic: requires-python
+Dynamic: summary
 
 Library that supports converting PyTorch models into a .tflite format, which can
 then be run with TensorFlow Lite and MediaPipe.  This enables applications for

{ai_edge_torch_nightly-0.7.0.dev20251017.dist-info → ai_edge_torch_nightly-0.7.0.dev20251018.dist-info}/RECORD

@@ -2,7 +2,7 @@ ai_edge_torch/__init__.py,sha256=lemyLCNoGYRnJsmDuGZu7qOqLbLqG6CGDFtu3ue1syU,129
 ai_edge_torch/_config.py,sha256=AiqhbcheF7j_ozIGDLC89k1we95aVgFDa-tR6h7UI0s,2529
 ai_edge_torch/conftest.py,sha256=r0GTrhMRhlmOGrrkvumHN8hkmyug6WvF60vWq8wRIBI,758
 ai_edge_torch/model.py,sha256=A7loFu8jE9CsXsfMmHYZ-KDFJiaD8Kkqwm_9d3IVzk0,5638
-ai_edge_torch/version.py,sha256=s37U1R9i7bjEl8Td8XZ5nmBvKsNMQV5vCFmKLpWXQ-k,806
+ai_edge_torch/version.py,sha256=07ZM69GTj0eqLoL4igeb-C1Ezp_B86je2hnvTkDH3Ho,806
 ai_edge_torch/_convert/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/_convert/conversion.py,sha256=iQk3R-pLq4c1nfLqPB4xTRj78gghxPGzJCJtILLdg5o,6123
 ai_edge_torch/_convert/conversion_utils.py,sha256=Sr8qXVcTwc-ZnZmK7yxVrIOOp1S_vNrwzC0zUvLTI2o,2160
@@ -250,11 +250,11 @@ ai_edge_torch/odml_torch/jax_bridge/__init__.py,sha256=e9Oa4J3An9FYr3zM0OzjzyNNi
 ai_edge_torch/odml_torch/jax_bridge/_wrap.py,sha256=LqwZ1vCJTSOzgzvH8LUAN-sAkF-l_pGj1AMEIzAqHCA,6638
 ai_edge_torch/odml_torch/jax_bridge/utils.py,sha256=T8isGc896VrHZ6c_L5pYmLpolQ7ibcOlgWfPuVFPzIg,2264
 ai_edge_torch/odml_torch/lowerings/__init__.py,sha256=uJ-niilt1c-D6QJzLwgvCUf62le_JsxQTlqj_iP_Ps0,1009
-ai_edge_torch/odml_torch/lowerings/_basic.py,sha256=sC4N5-7RS9yKecs97kM9J56enGvsZj1CJo7y79cuzRg,12784
+ai_edge_torch/odml_torch/lowerings/_basic.py,sha256=-zKZtqvOx3COBhjCDtiZWMn5fY-boktkWXjp5Kepiro,14716
 ai_edge_torch/odml_torch/lowerings/_batch_norm.py,sha256=PaLI0BB6pdBW1VyfW8VTOT_Be-ZcqYdNOsyfzKfq8Cg,2064
 ai_edge_torch/odml_torch/lowerings/_convolution.py,sha256=Q0aDzyUcZMoSzSbOU-r3LJMgPe6fble0QwdYVIOHHHk,6887
 ai_edge_torch/odml_torch/lowerings/_decomp_registry.py,sha256=ybOdoFE5HIJTkyiYcc73zpyUyUpioVnAca6k0wyJPs4,2572
-ai_edge_torch/odml_torch/lowerings/_jax_lowerings.py,sha256=6KZc0-SM0LupOnDAgnUjj47aHk7YJ0cC1ulAeWvV5_w,19289
+ai_edge_torch/odml_torch/lowerings/_jax_lowerings.py,sha256=7XlctokKu8jmXG51ZzdLz5HA7DDeD1bLai7aGUMs008,20457
 ai_edge_torch/odml_torch/lowerings/_layer_norm.py,sha256=khJIvDVk2s332Nd2Be-5dM6-wp5DGff61HCV5lskHmQ,3011
 ai_edge_torch/odml_torch/lowerings/_quantized_decomposed.py,sha256=XDZ0zLej_XaQDJnaAAxhNFAd7NfQm5SOVEp_nno_krA,6178
 ai_edge_torch/odml_torch/lowerings/_rand.py,sha256=g6SuqDkuC6hD35lyP1-5H7ASDIzPSmKukeNT5naZSv8,4133
@@ -270,8 +270,8 @@ ai_edge_torch/testing/__init__.py,sha256=_yGgvnBZWb7T3IN3mc4x1sS4vM96HZwM8pwIcPG
 ai_edge_torch/testing/export.py,sha256=k5mGDGzwc23Z4zaIVDs8CNh-oOt64gsf9MS9NjhbPy4,3293
 ai_edge_torch/testing/model_coverage/__init__.py,sha256=5P8J6Zk5YYtDvTBucFvB9NGSRI7Gw_24WnrbhXgycEE,765
 ai_edge_torch/testing/model_coverage/model_coverage.py,sha256=UPB448aMDUyC0HNYVqio2rcJPnDN0tBQMP08J6vPYew,4718
-ai_edge_torch_nightly-0.7.0.dev20251017.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.7.0.dev20251017.dist-info/METADATA,sha256=6izafx1yDIwodhe9vqegygTBcgLDDnXQuqtX1xFzwnY,2074
-ai_edge_torch_nightly-0.7.0.dev20251017.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-ai_edge_torch_nightly-0.7.0.dev20251017.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.7.0.dev20251017.dist-info/RECORD,,
+ai_edge_torch_nightly-0.7.0.dev20251018.dist-info/licenses/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.7.0.dev20251018.dist-info/METADATA,sha256=upHQ9ZpIm3SXQXAyfLoCWjW6fjiHFxZI5dfaRV_yZ8c,2297
+ai_edge_torch_nightly-0.7.0.dev20251018.dist-info/WHEEL,sha256=SmOxYU7pzNKBqASvQJ7DjX3XGUF92lrGhMb3R6_iiqI,91
+ai_edge_torch_nightly-0.7.0.dev20251018.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.7.0.dev20251018.dist-info/RECORD,,

{ai_edge_torch_nightly-0.7.0.dev20251017.dist-info → ai_edge_torch_nightly-0.7.0.dev20251018.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: bdist_wheel (0.45.1)
+Generator: setuptools (79.0.1)
 Root-Is-Purelib: true
 Tag: py3-none-any