We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Implement `binomial` and `beta` distribution functions in `keras.random`
Following up on the issue https://github.com/keras-team/keras/issues/18918

- Implement `binomial` and `beta` distribution functions in all backends currently supported by Keras namely TensorFlow, Jax, PyTorch and Numpy.
- Add unit tests for each of these functions

Importantly,
As tensorflow doesn't offer a built-in method for beta function so I've implemented a workaround using a statistical formula to use gamma distributed random variables to derive beta distributed random variable.
Specifically, $U(a, b) = X(a) / (X(a) + Y(b))$ where $U(a,b)$ is the beta distributed random variable using parameters $a$ and $b$ and $X(a)$ and $Y(b)$ are gamma-distributed random variables using parameter $a$ and $b$ respectively.

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `keras/backend/torch/random.py`
Content:
```
1 import torch
2 import torch._dynamo as dynamo
3 import torch.nn.functional as tnn
4 
5 from keras.backend.config import floatx
6 from keras.backend.torch.core import convert_to_tensor
7 from keras.backend.torch.core import get_device
8 from keras.backend.torch.core import to_torch_dtype
9 from keras.random.seed_generator import SeedGenerator
10 from keras.random.seed_generator import draw_seed
11 from keras.random.seed_generator import make_default_seed
12 
13 
14 # torch.Generator not supported with dynamo
15 # see: https://github.com/pytorch/pytorch/issues/88576
16 @dynamo.disable()
17 def torch_seed_generator(seed):
18     first_seed, second_seed = draw_seed(seed)
19     device = get_device()
20     if device == "meta":
21         # Generator is not supported by the meta device.
22         return None
23     generator = torch.Generator(device=get_device())
24     generator.manual_seed(int(first_seed + second_seed))
25     return generator
26 
27 
28 def normal(shape, mean=0.0, stddev=1.0, dtype=None, seed=None):
29     dtype = dtype or floatx()
30     dtype = to_torch_dtype(dtype)
31     # Do not use generator during symbolic execution.
32     if get_device() == "meta":
33         return torch.normal(
34             mean, stddev, size=shape, dtype=dtype, device=get_device()
35         )
36     generator = torch_seed_generator(seed)
37     return torch.normal(
38         mean,
39         stddev,
40         size=shape,
41         generator=generator,
42         dtype=dtype,
43         device=get_device(),
44     )
45 
46 
47 def categorical(logits, num_samples, dtype="int32", seed=None):
48     logits = convert_to_tensor(logits)
49     dtype = to_torch_dtype(dtype)
50     probs = torch.softmax(logits, dim=-1)
51     # Do not use generator during symbolic execution.
52     if get_device() == "meta":
53         return torch.multinomial(
54             probs,
55             num_samples,
56             replacement=True,
57         ).type(dtype)
58     generator = torch_seed_generator(seed)
59     return torch.multinomial(
60         probs,
61         num_samples,
62         replacement=True,
63         generator=generator,
64     ).type(dtype)
65 
66 
67 def uniform(shape, minval=0.0, maxval=1.0, dtype=None, seed=None):
68     dtype = dtype or floatx()
69     dtype = to_torch_dtype(dtype)
70     requested_shape = shape
71     if len(requested_shape) == 0:
72         shape = (1,)
73     # Do not use generator during symbolic execution.
74     if get_device() == "meta":
75         rand_tensor = torch.rand(size=shape, dtype=dtype, device=get_device())
76     else:
77         generator = torch_seed_generator(seed)
78         rand_tensor = torch.rand(
79             size=shape, generator=generator, dtype=dtype, device=get_device()
80         )
81 
82     output = (maxval - minval) * rand_tensor + minval
83 
84     if len(requested_shape) == 0:
85         return output[0]
86     return output
87 
88 
89 def randint(shape, minval, maxval, dtype="int32", seed=None):
90     dtype = to_torch_dtype(dtype)
91     # Do not use generator during symbolic execution.
92     if get_device() == "meta":
93         return torch.randint(
94             low=minval,
95             high=maxval,
96             size=shape,
97             dtype=dtype,
98             device=get_device(),
99         )
100     generator = torch_seed_generator(seed)
101     return torch.randint(
102         low=minval,
103         high=maxval,
104         size=shape,
105         generator=generator,
106         dtype=dtype,
107         device=get_device(),
108     )
109 
110 
111 def truncated_normal(shape, mean=0.0, stddev=1.0, dtype=None, seed=None):
112     # Take a larger standard normal dist, discard values outside 2 * stddev
113     # Offset by mean and stddev
114     x = normal(tuple(shape) + (4,), mean=0, stddev=1, dtype=dtype, seed=seed)
115     valid = (x > -2) & (x < 2)
116     indexes = valid.max(-1, keepdim=True)[1]
117     trunc_x = torch.empty(shape, device=get_device())
118     trunc_x.data.copy_(x.gather(-1, indexes).squeeze(-1))
119     trunc_x.data.mul_(stddev).add_(mean)
120     return trunc_x
121 
122 
123 def _get_concrete_noise_shape(inputs, noise_shape):
124     if noise_shape is None:
125         return inputs.shape
126 
127     concrete_inputs_shape = inputs.shape
128     concrete_noise_shape = []
129     for i, value in enumerate(noise_shape):
130         concrete_noise_shape.append(
131             concrete_inputs_shape[i] if value is None else value
132         )
133     return concrete_noise_shape
134 
135 
136 def dropout(inputs, rate, noise_shape=None, seed=None):
137     if (
138         seed is not None
139         and not (isinstance(seed, SeedGenerator) and seed._initial_seed is None)
140         or noise_shape is not None
141     ):
142         keep_prob = 1.0 - rate
143         noise_shape = _get_concrete_noise_shape(inputs, noise_shape)
144         keep_prob_matrix = torch.full(
145             noise_shape, keep_prob, device=get_device()
146         )
147         generator = torch_seed_generator(seed)
148 
149         # Do not use generator during symbolic execution.
150         if get_device() == "meta":
151             mask = torch.bernoulli(keep_prob_matrix)
152         else:
153             mask = torch.bernoulli(keep_prob_matrix, generator=generator)
154 
155         mask = mask.bool()
156         mask = torch.broadcast_to(mask, inputs.shape)
157         return torch.where(
158             mask,
159             inputs / keep_prob,
160             torch.zeros_like(inputs, dtype=inputs.dtype),
161         )
162     # Fast path, unseeded (since torch doesn't support seeding dropout!!!!)
163     # Using the above implementation is possible, but much slower.
164     return torch.nn.functional.dropout(
165         inputs, p=rate, training=True, inplace=False
166     )
167 
168 
169 def shuffle(x, axis=0, seed=None):
170     # Ref: https://github.com/pytorch/pytorch/issues/71409
171     x = convert_to_tensor(x)
172 
173     # Get permutation indices
174     # Do not use generator during symbolic execution.
175     if get_device() == "meta":
176         row_perm = torch.rand(x.shape[: axis + 1], device=get_device()).argsort(
177             axis
178         )
179     else:
180         generator = torch_seed_generator(seed)
181         row_perm = torch.rand(
182             x.shape[: axis + 1], generator=generator, device=get_device()
183         ).argsort(axis)
184     for _ in range(x.ndim - axis - 1):
185         row_perm.unsqueeze_(-1)
186 
187     # Reformat this for the gather operation
188     row_perm = row_perm.repeat(
189         *[1 for _ in range(axis + 1)], *(x.shape[axis + 1 :])
190     )
191     return x.gather(axis, row_perm)
192 
193 
194 def gamma(shape, alpha, dtype=None, seed=None):
195     dtype = dtype or floatx()
196     dtype = to_torch_dtype(dtype)
197     alpha = torch.ones(shape) * torch.tensor(alpha)
198     beta = torch.ones(shape)
199     prev_rng_state = torch.random.get_rng_state()
200     first_seed, second_seed = draw_seed(seed)
201     torch.manual_seed(first_seed + second_seed)
202     gamma_distribution = torch.distributions.gamma.Gamma(alpha, beta)
203     sample = gamma_distribution.sample().type(dtype)
204     torch.random.set_rng_state(prev_rng_state)
205     return sample
206 
207 
208 def binomial(shape, counts, probabilities, dtype=None, seed=None):
209     dtype = dtype or floatx()
210     dtype = to_torch_dtype(dtype)
211     counts = torch.ones(shape) * convert_to_tensor(counts)
212     probabilities = torch.ones(shape) * convert_to_tensor(probabilities)
213     prev_rng_state = torch.random.get_rng_state()
214     first_seed, second_seed = draw_seed(seed)
215     torch.manual_seed(first_seed + second_seed)
216     binomial_distribution = torch.distributions.binomial.Binomial(
217         total_count=counts, probs=probabilities
218     )
219     sample = binomial_distribution.sample().type(dtype)
220     torch.random.set_rng_state(prev_rng_state)
221     return sample
222 
223 
224 def beta(shape, alpha, beta, dtype=None, seed=None):
225     dtype = dtype or floatx()
226     dtype = to_torch_dtype(dtype)
227     alpha = torch.ones(shape) * convert_to_tensor(alpha)
228     beta = torch.ones(shape) * convert_to_tensor(beta)
229     prev_rng_state = torch.random.get_rng_state()
230     first_seed, second_seed = draw_seed(seed)
231     torch.manual_seed(first_seed + second_seed)
232     beta_distribution = torch.distributions.beta.Beta(
233         concentration1=alpha, concentration0=beta
234     )
235     sample = beta_distribution.sample().type(dtype)
236     torch.random.set_rng_state(prev_rng_state)
237     return sample
238
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/keras/backend/torch/random.py b/keras/backend/torch/random.py
--- a/keras/backend/torch/random.py
+++ b/keras/backend/torch/random.py
@@ -208,8 +208,8 @@
 def binomial(shape, counts, probabilities, dtype=None, seed=None):
     dtype = dtype or floatx()
     dtype = to_torch_dtype(dtype)
-    counts = torch.ones(shape) * convert_to_tensor(counts)
-    probabilities = torch.ones(shape) * convert_to_tensor(probabilities)
+    counts = torch.broadcast_to(convert_to_tensor(counts), shape)
+    probabilities = torch.broadcast_to(convert_to_tensor(probabilities), shape)
     prev_rng_state = torch.random.get_rng_state()
     first_seed, second_seed = draw_seed(seed)
     torch.manual_seed(first_seed + second_seed)
@@ -224,8 +224,8 @@
 def beta(shape, alpha, beta, dtype=None, seed=None):
     dtype = dtype or floatx()
     dtype = to_torch_dtype(dtype)
-    alpha = torch.ones(shape) * convert_to_tensor(alpha)
-    beta = torch.ones(shape) * convert_to_tensor(beta)
+    alpha = torch.broadcast_to(convert_to_tensor(alpha), shape)
+    beta = torch.broadcast_to(convert_to_tensor(beta), shape)
     prev_rng_state = torch.random.get_rng_state()
     first_seed, second_seed = draw_seed(seed)
     torch.manual_seed(first_seed + second_seed)

{"golden_diff": "diff --git a/keras/backend/torch/random.py b/keras/backend/torch/random.py\n--- a/keras/backend/torch/random.py\n+++ b/keras/backend/torch/random.py\n@@ -208,8 +208,8 @@\n def binomial(shape, counts, probabilities, dtype=None, seed=None):\n     dtype = dtype or floatx()\n     dtype = to_torch_dtype(dtype)\n-    counts = torch.ones(shape) * convert_to_tensor(counts)\n-    probabilities = torch.ones(shape) * convert_to_tensor(probabilities)\n+    counts = torch.broadcast_to(convert_to_tensor(counts), shape)\n+    probabilities = torch.broadcast_to(convert_to_tensor(probabilities), shape)\n     prev_rng_state = torch.random.get_rng_state()\n     first_seed, second_seed = draw_seed(seed)\n     torch.manual_seed(first_seed + second_seed)\n@@ -224,8 +224,8 @@\n def beta(shape, alpha, beta, dtype=None, seed=None):\n     dtype = dtype or floatx()\n     dtype = to_torch_dtype(dtype)\n-    alpha = torch.ones(shape) * convert_to_tensor(alpha)\n-    beta = torch.ones(shape) * convert_to_tensor(beta)\n+    alpha = torch.broadcast_to(convert_to_tensor(alpha), shape)\n+    beta = torch.broadcast_to(convert_to_tensor(beta), shape)\n     prev_rng_state = torch.random.get_rng_state()\n     first_seed, second_seed = draw_seed(seed)\n     torch.manual_seed(first_seed + second_seed)\n", "issue": "Implement `binomial` and `beta` distribution functions in `keras.random`\nFollowing up on the issue https://github.com/keras-team/keras/issues/18918\r\n\r\n- Implement `binomial` and `beta` distribution functions in all backends currently supported by Keras namely TensorFlow, Jax, PyTorch and Numpy.\r\n- Add unit tests for each of these functions\r\n\r\nImportantly,\r\nAs tensorflow doesn't offer a built-in method for beta function so I've implemented a workaround using a statistical formula to use gamma distributed random variables to derive beta distributed random variable.\r\nSpecifically, $U(a, b) = X(a) / (X(a) + Y(b))$ where $U(a,b)$ is the beta distributed random variable using parameters $a$ and $b$ and $X(a)$ and $Y(b)$ are gamma-distributed random variables using parameter $a$ and $b$ respectively.\n", "before_files": [{"content": "import torch\nimport torch._dynamo as dynamo\nimport torch.nn.functional as tnn\n\nfrom keras.backend.config import floatx\nfrom keras.backend.torch.core import convert_to_tensor\nfrom keras.backend.torch.core import get_device\nfrom keras.backend.torch.core import to_torch_dtype\nfrom keras.random.seed_generator import SeedGenerator\nfrom keras.random.seed_generator import draw_seed\nfrom keras.random.seed_generator import make_default_seed\n\n\n# torch.Generator not supported with dynamo\n# see: https://github.com/pytorch/pytorch/issues/88576\n@dynamo.disable()\ndef torch_seed_generator(seed):\n    first_seed, second_seed = draw_seed(seed)\n    device = get_device()\n    if device == \"meta\":\n        # Generator is not supported by the meta device.\n        return None\n    generator = torch.Generator(device=get_device())\n    generator.manual_seed(int(first_seed + second_seed))\n    return generator\n\n\ndef normal(shape, mean=0.0, stddev=1.0, dtype=None, seed=None):\n    dtype = dtype or floatx()\n    dtype = to_torch_dtype(dtype)\n    # Do not use generator during symbolic execution.\n    if get_device() == \"meta\":\n        return torch.normal(\n            mean, stddev, size=shape, dtype=dtype, device=get_device()\n        )\n    generator = torch_seed_generator(seed)\n    return torch.normal(\n        mean,\n        stddev,\n        size=shape,\n        generator=generator,\n        dtype=dtype,\n        device=get_device(),\n    )\n\n\ndef categorical(logits, num_samples, dtype=\"int32\", seed=None):\n    logits = convert_to_tensor(logits)\n    dtype = to_torch_dtype(dtype)\n    probs = torch.softmax(logits, dim=-1)\n    # Do not use generator during symbolic execution.\n    if get_device() == \"meta\":\n        return torch.multinomial(\n            probs,\n            num_samples,\n            replacement=True,\n        ).type(dtype)\n    generator = torch_seed_generator(seed)\n    return torch.multinomial(\n        probs,\n        num_samples,\n        replacement=True,\n        generator=generator,\n    ).type(dtype)\n\n\ndef uniform(shape, minval=0.0, maxval=1.0, dtype=None, seed=None):\n    dtype = dtype or floatx()\n    dtype = to_torch_dtype(dtype)\n    requested_shape = shape\n    if len(requested_shape) == 0:\n        shape = (1,)\n    # Do not use generator during symbolic execution.\n    if get_device() == \"meta\":\n        rand_tensor = torch.rand(size=shape, dtype=dtype, device=get_device())\n    else:\n        generator = torch_seed_generator(seed)\n        rand_tensor = torch.rand(\n            size=shape, generator=generator, dtype=dtype, device=get_device()\n        )\n\n    output = (maxval - minval) * rand_tensor + minval\n\n    if len(requested_shape) == 0:\n        return output[0]\n    return output\n\n\ndef randint(shape, minval, maxval, dtype=\"int32\", seed=None):\n    dtype = to_torch_dtype(dtype)\n    # Do not use generator during symbolic execution.\n    if get_device() == \"meta\":\n        return torch.randint(\n            low=minval,\n            high=maxval,\n            size=shape,\n            dtype=dtype,\n            device=get_device(),\n        )\n    generator = torch_seed_generator(seed)\n    return torch.randint(\n        low=minval,\n        high=maxval,\n        size=shape,\n        generator=generator,\n        dtype=dtype,\n        device=get_device(),\n    )\n\n\ndef truncated_normal(shape, mean=0.0, stddev=1.0, dtype=None, seed=None):\n    # Take a larger standard normal dist, discard values outside 2 * stddev\n    # Offset by mean and stddev\n    x = normal(tuple(shape) + (4,), mean=0, stddev=1, dtype=dtype, seed=seed)\n    valid = (x > -2) & (x < 2)\n    indexes = valid.max(-1, keepdim=True)[1]\n    trunc_x = torch.empty(shape, device=get_device())\n    trunc_x.data.copy_(x.gather(-1, indexes).squeeze(-1))\n    trunc_x.data.mul_(stddev).add_(mean)\n    return trunc_x\n\n\ndef _get_concrete_noise_shape(inputs, noise_shape):\n    if noise_shape is None:\n        return inputs.shape\n\n    concrete_inputs_shape = inputs.shape\n    concrete_noise_shape = []\n    for i, value in enumerate(noise_shape):\n        concrete_noise_shape.append(\n            concrete_inputs_shape[i] if value is None else value\n        )\n    return concrete_noise_shape\n\n\ndef dropout(inputs, rate, noise_shape=None, seed=None):\n    if (\n        seed is not None\n        and not (isinstance(seed, SeedGenerator) and seed._initial_seed is None)\n        or noise_shape is not None\n    ):\n        keep_prob = 1.0 - rate\n        noise_shape = _get_concrete_noise_shape(inputs, noise_shape)\n        keep_prob_matrix = torch.full(\n            noise_shape, keep_prob, device=get_device()\n        )\n        generator = torch_seed_generator(seed)\n\n        # Do not use generator during symbolic execution.\n        if get_device() == \"meta\":\n            mask = torch.bernoulli(keep_prob_matrix)\n        else:\n            mask = torch.bernoulli(keep_prob_matrix, generator=generator)\n\n        mask = mask.bool()\n        mask = torch.broadcast_to(mask, inputs.shape)\n        return torch.where(\n            mask,\n            inputs / keep_prob,\n            torch.zeros_like(inputs, dtype=inputs.dtype),\n        )\n    # Fast path, unseeded (since torch doesn't support seeding dropout!!!!)\n    # Using the above implementation is possible, but much slower.\n    return torch.nn.functional.dropout(\n        inputs, p=rate, training=True, inplace=False\n    )\n\n\ndef shuffle(x, axis=0, seed=None):\n    # Ref: https://github.com/pytorch/pytorch/issues/71409\n    x = convert_to_tensor(x)\n\n    # Get permutation indices\n    # Do not use generator during symbolic execution.\n    if get_device() == \"meta\":\n        row_perm = torch.rand(x.shape[: axis + 1], device=get_device()).argsort(\n            axis\n        )\n    else:\n        generator = torch_seed_generator(seed)\n        row_perm = torch.rand(\n            x.shape[: axis + 1], generator=generator, device=get_device()\n        ).argsort(axis)\n    for _ in range(x.ndim - axis - 1):\n        row_perm.unsqueeze_(-1)\n\n    # Reformat this for the gather operation\n    row_perm = row_perm.repeat(\n        *[1 for _ in range(axis + 1)], *(x.shape[axis + 1 :])\n    )\n    return x.gather(axis, row_perm)\n\n\ndef gamma(shape, alpha, dtype=None, seed=None):\n    dtype = dtype or floatx()\n    dtype = to_torch_dtype(dtype)\n    alpha = torch.ones(shape) * torch.tensor(alpha)\n    beta = torch.ones(shape)\n    prev_rng_state = torch.random.get_rng_state()\n    first_seed, second_seed = draw_seed(seed)\n    torch.manual_seed(first_seed + second_seed)\n    gamma_distribution = torch.distributions.gamma.Gamma(alpha, beta)\n    sample = gamma_distribution.sample().type(dtype)\n    torch.random.set_rng_state(prev_rng_state)\n    return sample\n\n\ndef binomial(shape, counts, probabilities, dtype=None, seed=None):\n    dtype = dtype or floatx()\n    dtype = to_torch_dtype(dtype)\n    counts = torch.ones(shape) * convert_to_tensor(counts)\n    probabilities = torch.ones(shape) * convert_to_tensor(probabilities)\n    prev_rng_state = torch.random.get_rng_state()\n    first_seed, second_seed = draw_seed(seed)\n    torch.manual_seed(first_seed + second_seed)\n    binomial_distribution = torch.distributions.binomial.Binomial(\n        total_count=counts, probs=probabilities\n    )\n    sample = binomial_distribution.sample().type(dtype)\n    torch.random.set_rng_state(prev_rng_state)\n    return sample\n\n\ndef beta(shape, alpha, beta, dtype=None, seed=None):\n    dtype = dtype or floatx()\n    dtype = to_torch_dtype(dtype)\n    alpha = torch.ones(shape) * convert_to_tensor(alpha)\n    beta = torch.ones(shape) * convert_to_tensor(beta)\n    prev_rng_state = torch.random.get_rng_state()\n    first_seed, second_seed = draw_seed(seed)\n    torch.manual_seed(first_seed + second_seed)\n    beta_distribution = torch.distributions.beta.Beta(\n        concentration1=alpha, concentration0=beta\n    )\n    sample = beta_distribution.sample().type(dtype)\n    torch.random.set_rng_state(prev_rng_state)\n    return sample\n", "path": "keras/backend/torch/random.py"}], "after_files": [{"content": "import torch\nimport torch._dynamo as dynamo\nimport torch.nn.functional as tnn\n\nfrom keras.backend.config import floatx\nfrom keras.backend.torch.core import convert_to_tensor\nfrom keras.backend.torch.core import get_device\nfrom keras.backend.torch.core import to_torch_dtype\nfrom keras.random.seed_generator import SeedGenerator\nfrom keras.random.seed_generator import draw_seed\nfrom keras.random.seed_generator import make_default_seed\n\n\n# torch.Generator not supported with dynamo\n# see: https://github.com/pytorch/pytorch/issues/88576\n@dynamo.disable()\ndef torch_seed_generator(seed):\n    first_seed, second_seed = draw_seed(seed)\n    device = get_device()\n    if device == \"meta\":\n        # Generator is not supported by the meta device.\n        return None\n    generator = torch.Generator(device=get_device())\n    generator.manual_seed(int(first_seed + second_seed))\n    return generator\n\n\ndef normal(shape, mean=0.0, stddev=1.0, dtype=None, seed=None):\n    dtype = dtype or floatx()\n    dtype = to_torch_dtype(dtype)\n    # Do not use generator during symbolic execution.\n    if get_device() == \"meta\":\n        return torch.normal(\n            mean, stddev, size=shape, dtype=dtype, device=get_device()\n        )\n    generator = torch_seed_generator(seed)\n    return torch.normal(\n        mean,\n        stddev,\n        size=shape,\n        generator=generator,\n        dtype=dtype,\n        device=get_device(),\n    )\n\n\ndef categorical(logits, num_samples, dtype=\"int32\", seed=None):\n    logits = convert_to_tensor(logits)\n    dtype = to_torch_dtype(dtype)\n    probs = torch.softmax(logits, dim=-1)\n    # Do not use generator during symbolic execution.\n    if get_device() == \"meta\":\n        return torch.multinomial(\n            probs,\n            num_samples,\n            replacement=True,\n        ).type(dtype)\n    generator = torch_seed_generator(seed)\n    return torch.multinomial(\n        probs,\n        num_samples,\n        replacement=True,\n        generator=generator,\n    ).type(dtype)\n\n\ndef uniform(shape, minval=0.0, maxval=1.0, dtype=None, seed=None):\n    dtype = dtype or floatx()\n    dtype = to_torch_dtype(dtype)\n    requested_shape = shape\n    if len(requested_shape) == 0:\n        shape = (1,)\n    # Do not use generator during symbolic execution.\n    if get_device() == \"meta\":\n        rand_tensor = torch.rand(size=shape, dtype=dtype, device=get_device())\n    else:\n        generator = torch_seed_generator(seed)\n        rand_tensor = torch.rand(\n            size=shape, generator=generator, dtype=dtype, device=get_device()\n        )\n\n    output = (maxval - minval) * rand_tensor + minval\n\n    if len(requested_shape) == 0:\n        return output[0]\n    return output\n\n\ndef randint(shape, minval, maxval, dtype=\"int32\", seed=None):\n    dtype = to_torch_dtype(dtype)\n    # Do not use generator during symbolic execution.\n    if get_device() == \"meta\":\n        return torch.randint(\n            low=minval,\n            high=maxval,\n            size=shape,\n            dtype=dtype,\n            device=get_device(),\n        )\n    generator = torch_seed_generator(seed)\n    return torch.randint(\n        low=minval,\n        high=maxval,\n        size=shape,\n        generator=generator,\n        dtype=dtype,\n        device=get_device(),\n    )\n\n\ndef truncated_normal(shape, mean=0.0, stddev=1.0, dtype=None, seed=None):\n    # Take a larger standard normal dist, discard values outside 2 * stddev\n    # Offset by mean and stddev\n    x = normal(tuple(shape) + (4,), mean=0, stddev=1, dtype=dtype, seed=seed)\n    valid = (x > -2) & (x < 2)\n    indexes = valid.max(-1, keepdim=True)[1]\n    trunc_x = torch.empty(shape, device=get_device())\n    trunc_x.data.copy_(x.gather(-1, indexes).squeeze(-1))\n    trunc_x.data.mul_(stddev).add_(mean)\n    return trunc_x\n\n\ndef _get_concrete_noise_shape(inputs, noise_shape):\n    if noise_shape is None:\n        return inputs.shape\n\n    concrete_inputs_shape = inputs.shape\n    concrete_noise_shape = []\n    for i, value in enumerate(noise_shape):\n        concrete_noise_shape.append(\n            concrete_inputs_shape[i] if value is None else value\n        )\n    return concrete_noise_shape\n\n\ndef dropout(inputs, rate, noise_shape=None, seed=None):\n    if (\n        seed is not None\n        and not (isinstance(seed, SeedGenerator) and seed._initial_seed is None)\n        or noise_shape is not None\n    ):\n        keep_prob = 1.0 - rate\n        noise_shape = _get_concrete_noise_shape(inputs, noise_shape)\n        keep_prob_matrix = torch.full(\n            noise_shape, keep_prob, device=get_device()\n        )\n        generator = torch_seed_generator(seed)\n\n        # Do not use generator during symbolic execution.\n        if get_device() == \"meta\":\n            mask = torch.bernoulli(keep_prob_matrix)\n        else:\n            mask = torch.bernoulli(keep_prob_matrix, generator=generator)\n\n        mask = mask.bool()\n        mask = torch.broadcast_to(mask, inputs.shape)\n        return torch.where(\n            mask,\n            inputs / keep_prob,\n            torch.zeros_like(inputs, dtype=inputs.dtype),\n        )\n    # Fast path, unseeded (since torch doesn't support seeding dropout!!!!)\n    # Using the above implementation is possible, but much slower.\n    return torch.nn.functional.dropout(\n        inputs, p=rate, training=True, inplace=False\n    )\n\n\ndef shuffle(x, axis=0, seed=None):\n    # Ref: https://github.com/pytorch/pytorch/issues/71409\n    x = convert_to_tensor(x)\n\n    # Get permutation indices\n    # Do not use generator during symbolic execution.\n    if get_device() == \"meta\":\n        row_perm = torch.rand(x.shape[: axis + 1], device=get_device()).argsort(\n            axis\n        )\n    else:\n        generator = torch_seed_generator(seed)\n        row_perm = torch.rand(\n            x.shape[: axis + 1], generator=generator, device=get_device()\n        ).argsort(axis)\n    for _ in range(x.ndim - axis - 1):\n        row_perm.unsqueeze_(-1)\n\n    # Reformat this for the gather operation\n    row_perm = row_perm.repeat(\n        *[1 for _ in range(axis + 1)], *(x.shape[axis + 1 :])\n    )\n    return x.gather(axis, row_perm)\n\n\ndef gamma(shape, alpha, dtype=None, seed=None):\n    dtype = dtype or floatx()\n    dtype = to_torch_dtype(dtype)\n    alpha = torch.ones(shape) * torch.tensor(alpha)\n    beta = torch.ones(shape)\n    prev_rng_state = torch.random.get_rng_state()\n    first_seed, second_seed = draw_seed(seed)\n    torch.manual_seed(first_seed + second_seed)\n    gamma_distribution = torch.distributions.gamma.Gamma(alpha, beta)\n    sample = gamma_distribution.sample().type(dtype)\n    torch.random.set_rng_state(prev_rng_state)\n    return sample\n\n\ndef binomial(shape, counts, probabilities, dtype=None, seed=None):\n    dtype = dtype or floatx()\n    dtype = to_torch_dtype(dtype)\n    counts = torch.broadcast_to(convert_to_tensor(counts), shape)\n    probabilities = torch.broadcast_to(convert_to_tensor(probabilities), shape)\n    prev_rng_state = torch.random.get_rng_state()\n    first_seed, second_seed = draw_seed(seed)\n    torch.manual_seed(first_seed + second_seed)\n    binomial_distribution = torch.distributions.binomial.Binomial(\n        total_count=counts, probs=probabilities\n    )\n    sample = binomial_distribution.sample().type(dtype)\n    torch.random.set_rng_state(prev_rng_state)\n    return sample\n\n\ndef beta(shape, alpha, beta, dtype=None, seed=None):\n    dtype = dtype or floatx()\n    dtype = to_torch_dtype(dtype)\n    alpha = torch.broadcast_to(convert_to_tensor(alpha), shape)\n    beta = torch.broadcast_to(convert_to_tensor(beta), shape)\n    prev_rng_state = torch.random.get_rng_state()\n    first_seed, second_seed = draw_seed(seed)\n    torch.manual_seed(first_seed + second_seed)\n    beta_distribution = torch.distributions.beta.Beta(\n        concentration1=alpha, concentration0=beta\n    )\n    sample = beta_distribution.sample().type(dtype)\n    torch.random.set_rng_state(prev_rng_state)\n    return sample\n", "path": "keras/backend/torch/random.py"}]}