{"id":13688752,"url":"https://github.com/JiaYaobo/fenbux","last_synced_at":"2025-05-01T20:30:38.374Z","repository":{"id":179062257,"uuid":"662644589","full_name":"JiaYaobo/fenbux","owner":"JiaYaobo","description":"A Simple Statistical Distribution Library in JAX","archived":false,"fork":false,"pushed_at":"2024-03-30T07:10:41.000Z","size":819,"stargazers_count":16,"open_issues_count":5,"forks_count":0,"subscribers_count":2,"default_branch":"main","last_synced_at":"2024-11-12T12:48:39.764Z","etag":null,"topics":["jax","probabilistic-programming","statistical-learning"],"latest_commit_sha":null,"homepage":"https://jiayaobo.github.io/fenbux/","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"apache-2.0","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/JiaYaobo.png","metadata":{"files":{"readme":"readme.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null}},"created_at":"2023-07-05T15:14:52.000Z","updated_at":"2024-02-12T20:26:33.000Z","dependencies_parsed_at":"2024-01-06T10:25:14.462Z","dependency_job_id":"944ff4cc-7a7b-446f-9798-7d4dab084c34","html_url":"https://github.com/JiaYaobo/fenbux","commit_stats":null,"previous_names":["jiayaobo/fenbux"],"tags_count":2,"template":false,"template_full_name":null,"repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/JiaYaobo%2Ffenbux","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/JiaYaobo%2Ffenbux/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/JiaYaobo%2Ffenbux/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/JiaYaobo%2Ffenbux/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/JiaYaobo","download_url":"https://codeload.github.com/JiaYaobo/fenbux/tar.gz/refs/heads/main","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":251940419,"owners_count":21668531,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["jax","probabilistic-programming","statistical-learning"],"created_at":"2024-08-02T15:01:21.753Z","updated_at":"2025-05-01T20:30:38.369Z","avatar_url":"https://github.com/JiaYaobo.png","language":"Python","funding_links":[],"categories":["Python"],"sub_categories":[],"readme":"# FenbuX\n\n*A Simple Probalistic Distribution Library in JAX*\n\n*fenbu* (分布, pronounce like: /fen'bu:/)-X is a simple probalistic distribution library in JAX. In fenbux, We provide you:\n\n* A simple and easy-to-use interface like **Distributions.jl**\n* Bijectors like **TensorFlow Probability** and **Bijector.jl**\n* PyTree input/output\n* Multiple dispatch for different distributions based on [plum-dispatch](https://github.com/beartype/plum)\n* All jax feautures (vmap, pmap, jit, autograd etc.)\n\nSee [document](https://jiayaobo.github.io/fenbux/)\n\n## Examples\n\n### Statistics of Distributions 🤔\n\n```python\nimport jax.numpy as jnp\nfrom fenbux import variance, skewness, mean\nfrom fenbux.univariate import Normal\n\nμ = {'a': jnp.array([1., 2., 3.]), 'b': jnp.array([4., 5., 6.])} \nσ = {'a': jnp.array([4., 5., 6.]), 'b': jnp.array([7., 8., 9.])}\n\ndist = Normal(μ, σ)\nmean(dist) # {'a': Array([1., 2., 3.], dtype=float32), 'b': Array([4., 5., 6.], dtype=float32)}\nvariance(dist) # {'a': Array([16., 25., 36.], dtype=float32), 'b': Array([49., 64., 81.], dtype=float32)}\nskewness(dist) # {'a': Array([0., 0., 0.], dtype=float32), 'b': Array([0., 0., 0.], dtype=float32)}\n```\n\n### Random Variables Generation\n\n```python\nimport jax.random as jr\nfrom fenbux import rand\nfrom fenbux.univariate import Normal\n\n\nkey =  jr.PRNGKey(0)\nx = {'a': {'c': {'d': {'e': 1.}}}}\ny = {'a': {'c': {'d': {'e': 1.}}}}\n\ndist = Normal(x, y)\nrand(dist, key, shape=(3, )) # {'a': {'c': {'d': {'e': Array([1.6248107 , 0.69599575, 0.10169095], dtype=float32)}}}}\n```\n\n### Evaluations of Distribution 👩‍🎓\n\nCDF, PDF, and more...\n\n```python\nimport jax.numpy as jnp\nfrom fenbux import cdf, logpdf\nfrom fenbux.univariate import Normal\n\n\nμ = jnp.array([1., 2., 3.])\nσ = jnp.array([4., 5., 6.])\n\ndist = Normal(μ, σ)\ncdf(dist, jnp.array([1., 2., 3.])) # Array([0.5, 0.5, 0.5], dtype=float32)\nlogpdf(dist, jnp.array([1., 2., 3.])) # Array([-2.305233 , -2.5283763, -2.7106981], dtype=float32)\n```\n\n### Nested Transformations of Distribution 🤖\n\n```python\nimport fenbux as fbx\nimport jax.numpy as jnp\nfrom fenbux.univariate import Normal\n\n# truncate and censor and affine\nd = Normal(0, 1)\nfbx.affine(fbx.censor(fbx.truncate(d, 0, 1), 0, 1), 0, 1)\nfbx.logpdf(d, 0.5)\n```\n\n```\nArray(-1.0439385, dtype=float32)\n```\n\n### Compatible with JAX transformations 😃\n\n- vmap\n\n```python\nimport jax.numpy as jnp\nfrom jax import vmap\n\nfrom fenbux import logpdf\nfrom fenbux.univariate import Normal\n\n\ndist = Normal({'a': jnp.zeros((2, 3))}, {'a':jnp.ones((2, 3, 5))}) # each batch shape is (2, 3)\nx = jnp.zeros((2, 3, 5))\n# claim use_batch=True to use vmap\nvmap(logpdf, in_axes=(Normal(None, {'a': 2}, use_batch=True), 2))(dist, x) \n```\n\n- grad\n\n```python\nimport jax.numpy as jnp\nfrom jax import jit, grad\nfrom fenbux import logpdf\nfrom fenbux.univariate import Normal\n\ndist = Normal(0., 1.)\ngrad(logpdf)(dist, 0.)\n```\n\n### Bijectors 🧙‍♂️\n\nEvaluate a bijector\n\n```python\nimport jax.numpy as jnp\nfrom fenbux.bijector import Exp, evaluate\n\nbij = Exp()\nx = jnp.array([1., 2., 3.])\n\nevaluate(bij, x)\n```\n\nApply a bijector to a distribution\n\n```python\nimport jax.numpy as jnp\nfrom fenbux.bijector import Exp, transform\nfrom fenbux.univariate import Normal\nfrom fenbux import logpdf\n\ndist = Normal(0, 1)\nbij = Exp()\n\nlog_normal = transform(dist, bij)\n\nx = jnp.array([1., 2., 3.])\nlogpdf(log_normal, x)\n```\n\n### Speed 🔦\n\n* Common Evaluations\n\n```python\nimport numpy as np\nfrom scipy.stats import norm\nfrom jax import jit\nfrom fenbux import logpdf, rand\nfrom fenbux.univariate import Normal\nfrom tensorflow_probability.substrates.jax.distributions import Normal as Normal2\n\ndist = Normal(0, 1)\ndist2 = Normal2(0, 1)\ndist3 = norm(0, 1)\nx = np.random.normal(size=100000)\n\n%timeit jit(logpdf)(dist, x).block_until_ready()\n%timeit jit(dist2.log_prob)(x).block_until_ready()\n%timeit dist3.logpdf(x)\n```\n\n```\n51.2 µs ± 1.47 µs per loop (mean ± std. dev. of 7 runs, 10,000 loops each)\n11.1 ms ± 176 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n1.12 ms ± 20.1 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)\n```\n\n* Evaluations with Bijector Transformed Distributions\n\n```python\nimport jax.numpy as jnp\nimport numpy as np\nimport tensorflow_probability.substrates.jax.bijectors as tfb\nimport tensorflow_probability.substrates.jax.distributions as tfd\nfrom jax import jit\n\nfrom fenbux import logpdf\nfrom fenbux.bijector import Exp, transform\nfrom fenbux.univariate import Normal\n\n\nx = jnp.asarray(np.random.uniform(size=100000))\ndist = Normal(0, 1)\nbij = Exp()\nlog_normal = transform(dist, bij)\n\ndist2 = tfd.Normal(loc=0, scale=1)\nbij2 = tfb.Exp()\nlog_normal2 = tfd.TransformedDistribution(dist2, bij2)\n\ndef log_prob(d, x):\n    return d.log_prob(x)\n\n%timeit jit(logpdf)(log_normal, x).block_until_ready()\n%timeit jit(log_prob)(log_normal2, x).block_until_ready()\n```\n\n```\n131 µs ± 514 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)\n375 µs ± 10.9 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)\n```\n\n\n## Installation\n\n* Install on your local device.\n\n```bash\ngit clone https://github.com/JiaYaobo/fenbux.git\npip install -e .\n```\n\n* Install from PyPI.\n\n```bash\npip install -U fenbux\n```\n\n## Reference\n\n* [Distributions.jl](https://github.com/JuliaStats/Distributions.jl)\n* [Equinox](https://github.com/patrick-kidger/equinox)\n\n\n## Citation\n\n```bibtex\n@software{fenbux,\n  author = {Jia, Yaobo},\n  title = {fenbux: A Simple Probalistic Distribution Library in JAX},\n  url = {https://github.com/JiaYaobo/fenbux},\n  year = {2024}\n}\n```\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2FJiaYaobo%2Ffenbux","html_url":"https://awesome.ecosyste.ms/projects/github.com%2FJiaYaobo%2Ffenbux","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2FJiaYaobo%2Ffenbux/lists"}