jax.random.categorical doesn't respect shape parameter with broadcasted logits
### Description

The docs for `jax.random.categorical` say that the output of `jax.random.categorical` is `shape` when `shape` is provided. However, when provided `logits` which have to be broadcasted with `shape`, it just ignores the shape parameter after checking that it's broadcast compatible. I believe this can be fixed by broadcasting logits appropriately.

```
import jax
out = jax.random.categorical(jax.random.PRNGKey(0), logits=jax.numpy.ones( (4, 1, 10)), axis=-1, shape=(4, 8))

# out.shape is (4, 1), should be (4, 8)

assert out.shape == (4, 8), f"out.shape is {out.shape}, should be (4, 8)"
```



### What jax/jaxlib version are you using?

0.3.23

### Which accelerator(s) are you using?

CPU

### Additional system info

_No response_

### NVIDIA GPU info

_No response_
diff --git a/jax/_src/random.py b/jax/_src/random.py
index 0e5e01dce..1292c8350 100644
--- a/jax/_src/random.py
+++ b/jax/_src/random.py
@@ -1327,10 +1327,12 @@ def categorical(key: KeyArray,
     shape = tuple(shape)
     _check_shape("categorical", shape, batch_shape)
 
-  sample_shape = shape[:len(shape)-len(batch_shape)]
+  shape_prefix = shape[:len(shape)-len(batch_shape)]
+  logits_shape = list(shape[len(shape) - len(batch_shape):])
+  logits_shape.insert(axis % len(logits_arr.shape), logits_arr.shape[axis])
   return jnp.argmax(
-      gumbel(key, sample_shape + logits_arr.shape, logits_arr.dtype) +
-      lax.expand_dims(logits_arr, tuple(range(len(sample_shape)))),
+      gumbel(key, (*shape_prefix, *logits_shape), logits_arr.dtype) +
+      lax.expand_dims(logits_arr, tuple(range(len(shape_prefix)))),
       axis=axis)
 
 
diff --git a/tests/random_test.py b/tests/random_test.py
index 5cebbfd63..20802e990 100644
--- a/tests/random_test.py
+++ b/tests/random_test.py
@@ -1429,6 +1429,19 @@ class LaxRandomTest(jtu.JaxTestCase):
     # just lower, don't run, takes too long
     jax.jit(f).lower()
 
+  @jtu.sample_product(shape=[(3, 4)],
+                      logits_shape_base=[(3, 4), (3, 1), (1, 4)],
+                      axis=[-3, -2, -1, 0, 1, 2])
+  def test_categorical_shape_argument(self, shape, logits_shape_base, axis):
+    # https://github.com/google/jax/issues/13124
+    logits_shape = list(logits_shape_base)
+    logits_shape.insert(axis % (len(logits_shape_base) + 1), 10)
+    assert logits_shape[axis] == 10
+    logits = jnp.ones(logits_shape)
+    samples = jax.random.categorical(jax.random.PRNGKey(0), logits=logits,
+                                     axis=axis, shape=shape)
+    self.assertEqual(samples.shape, shape)
+
 
 class KeyArrayTest(jtu.JaxTestCase):
   # Key arrays involve:
['jax/_src/random.py', 'tests/random_test.py']