Fix JAX multinomial size handling

[FlorianPfaff]

Summary

• add NumPy-style size= support to the JAX random backend's multinomial
• validate multinomial trial counts and probability vector shape before sampling
• add JAX-specific regression tests for multidraw output shape, invalid sizes, and invalid trial counts

Validation

• Exercised the new JAX multinomial sampling logic locally in isolation with JAX 0.9.0.1 for size=None, size=(), scalar size, tuple size, and zero-sized samples.
• Full repository checkout/tests could not be run from the sandbox because direct GitHub cloning failed with DNS resolution for github.com; CI should cover the integrated test matrix on the PR.

[github-actions]

✅MegaLinter analysis: Success

Descriptor	Linter	Files	Fixed	Errors	Warnings	Elapsed time
✅ COPYPASTE	jscpd	yes		no	no	46.89s
✅ JSON	prettier	7	0	0	0	1.04s
✅ JSON	v8r	7		0	0	4.63s
✅ MARKDOWN	markdownlint	66	0	0	0	1.78s
✅ MARKDOWN	markdown-table-formatter	66	0	0	0	0.81s
✅ PYTHON	black	786	1	0	0	38.06s
✅ PYTHON	isort	786	1	0	0	2.29s
✅ REPOSITORY	checkov	yes		no	no	46.88s
✅ REPOSITORY	gitleaks	yes		no	no	17.48s
✅ REPOSITORY	git_diff	yes		no	no	0.12s
✅ REPOSITORY	secretlint	yes		no	no	22.5s
✅ REPOSITORY	syft	yes		no	no	4.87s
✅ REPOSITORY	trivy-sbom	yes		no	no	6.76s
✅ REPOSITORY	trufflehog	yes		no	no	22.66s
✅ YAML	prettier	11	0	0	0	0.75s
✅ YAML	v8r	11		0	0	11.23s
✅ YAML	yamllint	11		0	0	0.52s

Notices

📣 MegaLinter 9.5.0 is out! Discover the new features and security recommendations in the release announcement. (Skip this info by defining SECURITY_SUGGESTIONS: false)

See detailed reports in MegaLinter artifacts

Your project could benefit from a custom flavor, which would allow you to run only the linters you need, and thus improve runtime performances. (Skip this info by defining FLAVOR_SUGGESTIONS: false)

• Documentation: Custom Flavors
• Command: npx mega-linter-runner@9.5.0 --custom-flavor-setup --custom-flavor-linters PYTHON_BLACK,PYTHON_ISORT,COPYPASTE_JSCPD,JSON_V8R,JSON_PRETTIER,MARKDOWN_MARKDOWNLINT,MARKDOWN_MARKDOWN_TABLE_FORMATTER,REPOSITORY_CHECKOV,REPOSITORY_GIT_DIFF,REPOSITORY_GITLEAKS,REPOSITORY_SECRETLINT,REPOSITORY_SYFT,REPOSITORY_TRIVY_SBOM,REPOSITORY_TRUFFLEHOG,YAML_PRETTIER,YAML_YAMLLINT,YAML_V8R

Show us your support by starring ⭐ the repository