returnn 1.20250814.205214__tar.gz → 1.20250817.33823__tar.gz

{returnn-1.20250814.205214/returnn.egg-info → returnn-1.20250817.33823}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: returnn
-Version: 1.20250814.205214
+Version: 1.20250817.33823
 Summary: The RWTH extensible training framework for universal recurrent neural networks
 Home-page: https://github.com/rwth-i6/returnn/
 Author: Albert Zeyer

returnn-1.20250817.33823/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20250817.033823'
+long_version = '1.20250817.033823+git.4f41ff3'

{returnn-1.20250814.205214 → returnn-1.20250817.33823}/returnn/frontend/_backend.py

@@ -1153,6 +1153,13 @@ class Backend(Generic[T]):
         """
         raise NotImplementedError
 
+    @staticmethod
+    def random_choice_with_replacement(dims: Sequence[Dim], *, probs: Tensor, axis: Dim) -> Tensor:
+        """
+        random choice with replacement. See `rf.random_choice_with_replacement` for details.
+        """
+        raise NotImplementedError
+
     @staticmethod
     def masked_select(
         tensor: Tensor, *, mask: Tensor, dims: Sequence[Dim], out_dim: Optional[Dim] = None

{returnn-1.20250814.205214 → returnn-1.20250817.33823}/returnn/frontend/rand.py

@@ -65,6 +65,7 @@ __all__ = [
     "random_normal",
     "random_truncated_normal",
     "random_choice_without_replacement",
+    "random_choice_with_replacement",
 ]
 
 
@@ -379,3 +380,16 @@ def random_choice_without_replacement(
     scores = log_probs + scores_random_sample
     _, indices, _ = rf.top_k(scores, k_dim=num_samples_dim, axis=axis)
     return indices
+
+
+def random_choice_with_replacement(dims: Sequence[Dim], *, probs: Tensor, axis: Dim) -> Tensor:
+    """
+    Randomly sample with replacement.
+
+    :param dims: {common_dims..., new_dims...}. Defines how many samples to draw. Defines the output shape.
+    :param probs: {common_dims..., axis}
+    :param axis: feature axis, where to sample from
+    :return: random indices shape dims -> axis.
+    """
+    # noinspection PyProtectedMember
+    return probs._raw_backend.random_choice_with_replacement(dims=dims, probs=probs, axis=axis)

{returnn-1.20250814.205214 → returnn-1.20250817.33823}/returnn/torch/frontend/_backend.py

@@ -1718,6 +1718,28 @@ class TorchBackend(Backend[torch.Tensor]):
             )
         return out
 
+    @staticmethod
+    def random_choice_with_replacement(dims: Sequence[Dim], *, probs: Tensor, axis: Dim) -> Tensor:
+        """random choice with replacement"""
+        assert all(d == axis or d in dims for d in probs.dims), (
+            f"random_choice_with_replacement: dims {dims} not compatible with probs {probs} and axis {axis}"
+        )
+        common_dims = [d for d in dims if d in probs.dims]
+        assert axis not in common_dims
+        probs = probs.copy_transpose(common_dims + [axis])
+        non_common_dims = [d for d in dims if d not in common_dims]
+        num_samples = prod([d.get_dim_value() for d in non_common_dims])
+        if len(common_dims) >= 2:
+            probs, flat_common_dim = rf.merge_dims(probs, dims=common_dims)
+        out_raw = torch.multinomial(probs.raw_tensor, num_samples=num_samples, replacement=True)
+        out_raw = out_raw.reshape(
+            [d.get_dim_value() for d in common_dims] + [d.get_dim_value() for d in non_common_dims]
+        )
+        out = rf.convert_to_tensor(out_raw, dims=common_dims + non_common_dims, sparse_dim=axis)
+        out = out.copy_transpose(dims)
+        out.name = "random_choice_with_replacement"
+        return out
+
     @staticmethod
     def masked_select(
         tensor: Tensor, *, mask: Tensor, dims: Sequence[Dim], out_dim: Optional[Dim] = None

{returnn-1.20250814.205214 → returnn-1.20250817.33823/returnn.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: returnn
-Version: 1.20250814.205214
+Version: 1.20250817.33823
 Summary: The RWTH extensible training framework for universal recurrent neural networks
 Home-page: https://github.com/rwth-i6/returnn/
 Author: Albert Zeyer

{returnn-1.20250814.205214 → returnn-1.20250817.33823}/tests/test_rf_base.py

@@ -861,3 +861,48 @@ def test_copy_compatible_to_empty():
             torch.testing.assert_close(out_raw, out.raw_tensor)
 
     run_model(extern_data, lambda **_kwargs: rf.Module(), _forward_step, test_tensorflow=False)
+
+
+def test_random_uniform_int():
+    batch_dim = Dim(3, name="batch")
+    feat_dim = Dim(5, name="feat")
+
+    # noinspection PyShadowingNames
+    def _forward_step(**_kwargs):
+        out = rf.random_uniform([batch_dim, feat_dim], dtype="int32", minval=-10, maxval=10)
+        out.mark_as_default_output(shape=[batch_dim, feat_dim])
+
+    run_model(TensorDict(), lambda **_kwargs: rf.Module(), _forward_step)
+
+
+def test_random_choice_with_replacement():
+    batch_dim = Dim(3, name="batch")
+    time_dim = Dim(7, name="time")
+    feat_dim = Dim(5, name="feat")
+
+    # noinspection PyShadowingNames
+    def _forward_step(**_kwargs):
+        probs = rf.softmax(rf.random_uniform([feat_dim], dtype="float32", minval=-1.0, maxval=1.0), axis=feat_dim)
+        out = rf.random_choice_with_replacement([batch_dim], probs=probs, axis=feat_dim)
+        assert out.dims == (batch_dim,) and out.sparse_dim == feat_dim
+        out.mark_as_default_output(shape=[batch_dim])
+
+        probs = rf.softmax(
+            rf.random_uniform([batch_dim, feat_dim], dtype="float32", minval=-1.0, maxval=1.0), axis=feat_dim
+        )
+        out = rf.random_choice_with_replacement([batch_dim], probs=probs, axis=feat_dim)
+        assert out.dims == (batch_dim,) and out.sparse_dim == feat_dim
+        out.mark_as_output("batched_no_new", shape=[batch_dim])
+
+        out = rf.random_choice_with_replacement([batch_dim, time_dim], probs=probs, axis=feat_dim)
+        assert out.dims == (batch_dim, time_dim) and out.sparse_dim == feat_dim
+        out.mark_as_output("batched_time", shape=[batch_dim, time_dim])
+
+        probs = rf.softmax(
+            rf.random_uniform([batch_dim, time_dim, feat_dim], dtype="float32", minval=-1.0, maxval=1.0), axis=feat_dim
+        )
+        out = rf.random_choice_with_replacement([batch_dim, time_dim], probs=probs, axis=feat_dim)
+        assert out.dims == (batch_dim, time_dim) and out.sparse_dim == feat_dim
+        out.mark_as_output("mult_common", shape=[batch_dim, time_dim])
+
+    run_model(TensorDict(), lambda **_kwargs: rf.Module(), _forward_step, test_tensorflow=False)

returnn-1.20250814.205214/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20250814.205214'
-long_version = '1.20250814.205214+git.26a4c87'