Direct sampling & plots

Sample distribution

Every distribution in m.dist can be sampled directly, outside of a model, by passing sample=True. This is useful for prior predictive checks, simulating data, or quickly inspecting a distribution’s shape.

from BayesForge import bf
m = bf()
# Draw 1000 samples from a standard normal
m.dist.normal(0, 1, shape=(4,), sample=True)

/home/sosa/work/3.12venv/lib/python3.13/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html
  from .autonotebook import tqdm as notebook_tqdm
An NVIDIA GPU may be present on this machine, but a CUDA-enabled jaxlib is not installed. Falling back to cpu.

bf v 0.0.58 package loaded
jax.local_device_count 32

Array([-0.5550907 ,  0.37154715, -0.26478447,  0.13819478], dtype=float64)

By using (to_jax=False), a direct sample is returned as a SampledData object — a thin wrapper around a JAX array. It behaves like a JAX array everywhere (arithmetic, indexing, jnp.* functions, jit/vmap/grad, use as a distribution parameter), and its repr reports the dtype so you can see the kind and precision at a glance:

The extra thing a SampledData gives you over a raw array is a set of built-in plotting methods and built-in helper functions.

Plotting helpers

Each plotting method is called directly on the sampled object. Most accept interactive=True (the default, Plotly) or interactive=False (matplotlib / seaborn):

The available methods depend on the shape of the sampled array.

1D-(n_samples,)

histogram

m.dist.normal(0, 1, shape=(1000,), sample=True, to_jax=False).hist() 

kernel

m.dist.normal(0, 1, shape=(1000,), sample=True, to_jax=False).density()

PPC

m.dist.normal(0,1,shape=(100,4),sample=True, to_jax=False).ppc_plot(m.dist.normal(0,1,shape=(100,),sample=True, to_jax=False))

Time series — (n_samples, n_time)

Treat the second axis as time and summarise across draws.

m.dist.gaussian_random_walk(scale=0.5, num_steps=50, shape=(500,), sample=True, to_jax=False).timeseries(credible_interval=0.9)

2D-(n_samples, n_variables)

histogram

m.dist.normal(0, 1, shape=(1000,4), sample=True, to_jax=False).hist() 

kernel

m.dist.normal(0, 1, shape=(1000,4), sample=True, to_jax=False).density() 

boxplot

m.dist.normal(0, 1, shape=(1000,4), sample=True, to_jax=False).boxplot() 

violinplot

m.dist.normal(0, 1, shape=(1000,4), sample=True, to_jax=False).violinplot() 

pairplot

m.dist.normal(0, 1, shape=(1000,4), sample=True, to_jax=False).pairplot() 

corr_heatmap

m.dist.lkj_cholesky(4, 2, name = 'a', sample = True, to_jax=False).corr_heatmap() 

traceplot

m.dist.normal(0, 1, shape=(1000,4), sample=True, to_jax=False).traceplot() 

surface_3d

Computes a 2D Kernel Density Estimate (KDE) using the first two variables. It plots a 3D surface where the X and Y axes represent the variables and the Z axis represents the estimated probability density.

m.dist.normal(0, 1, shape=(1000,4), sample=True, to_jax=False).surface_3d() 

scatter3d

Takes the first three variables (columns) and plots them directly as points in a 3D space (X, Y, and Z axes). Requires the array to have at least 3 variables (columns).

m.dist.normal(0, 1, shape=(1000,4), sample=True, to_jax=False).scatter3d() 

Posterior predictive

Compare draws against observed data.

3D-(n_samples, n_groups, n_times)

3D scatter — (n_samples, >=3)

m.dist.normal(0,1, name = 'a', shape=(100,4, 5), sample = True, to_jax=False).hist()

Non-plot helper

In addition to plot helper functions, there are several other helper functions that can be used to compute summaries or transform the sampled array. For reduction operations (like mean, std, sum), you can pass an axis argument to reduce along specific dimensions. For example, passing axis=0 will reduce across the sample dimension.

mean

m.dist.normal(0,1, name = 'a', shape=(100,4), sample = True, to_jax=False).mean(axis=0)

SampledData([-0.13334232,  0.2501823 ,  0.00978141,  0.06576618], dtype=float64)

std

m.dist.normal(0,1, name = 'a', shape=(100,4), sample = True, to_jax=False).std()

Array(1.02588659, dtype=float64)

variance

m.dist.normal(0,1, name = 'a', shape=(100,4), sample = True, to_jax=False).var()

Array(0.96947909, dtype=float64)

hdi

m.dist.normal(0,1, name = 'a', shape=(100,4), sample = True, to_jax=False).hdi(cred_mass=0.95)

(Array([-0.07487499,  0.40088192,  0.40251631,  2.18534984], dtype=float64),
 Array([-0.45834097,  0.17091522,  1.276105  ,  1.32610808], dtype=float64))

sum

m.dist.normal(0,1, name = 'a', shape=(100,4), sample = True, to_jax=False).sum()

Array(-35.21211041, dtype=float64)

prod

m.dist.normal(0,1, name = 'a', shape=(100,4), sample = True, to_jax=False).prod()

Array(1.31092218e-98, dtype=float64)

min

m.dist.normal(0,1, name = 'a', shape=(100,4), sample = True, to_jax=False).min()

Array(-2.68929725, dtype=float64)

max

m.dist.normal(0,1, name = 'a', shape=(100,4, 5), sample = True, to_jax=False).max()

Array(3.42966119, dtype=float64)

argmin

m.dist.normal(0,1, name = 'a', shape=(100,4), sample = True, to_jax=False).argmin()

Array(84, dtype=int64)

argmax

m.dist.normal(0,1, name = 'a', shape=(100,4), sample = True, to_jax=False).argmax()

Array(233, dtype=int64)

all

m.dist.normal(0,1, name = 'a', shape=(100,4, 5), sample = True, to_jax=False).all()

Array(True, dtype=bool)

any

m.dist.normal(0,1, name = 'a', shape=(100,4), sample = True, to_jax=False).any()

Array(True, dtype=bool)

cumsum

m.dist.normal(0,1, name = 'a', shape=(100,4), sample = True, to_jax=False).cumsum()

SampledData([-0.63578926, -0.17804941, -1.58290785, -1.87170156, -4.14525923,
       -3.11694519, -4.42013083, -4.37601634, -3.07856047, -2.04925911,
       -3.03124515, -2.45998945, -2.3847695 , -2.97060164, -4.66621447,
       -4.96842123, -4.22842776, -5.10248091, -4.52046884, -5.3060986 ,
       -4.64257205, -4.95388316, -3.01393981, -3.77209503, -4.39617448,
       -4.84646193, -4.6616556 , -6.14916896, -6.42734398, -6.67100089,
       -6.07411741, -6.66249441, -5.36782531, -6.69846924, -6.21621535,
       -5.31175873, -3.49531165, -4.12087376, -3.6987261 , -3.06733673,
       -3.61661148, -3.73978088, -3.28535465, -2.99780189, -3.31704086,
       -2.14784496, -0.24828819,  0.30881405, -0.53899346, -1.34413752,
       -1.34584539, -2.25573194, -3.18784912, -2.68146834, -2.11895077,
       -0.41203008, -0.67130641, -1.52312884, -1.33449293, -1.57481445,
       -0.10058   , -1.30422323, -1.78878153, -2.35574265, -2.47375093,
       -1.3478823 , -2.45293663, -2.2221739 , -2.81437199, -3.27089587,
       -2.50402216, -3.46001596, -2.9674292 , -2.72351172, -1.93675804,
       -1.22707337, -0.69001749, -1.59696727, -2.36020205, -2.64863803,
       -3.24146337, -3.1839793 , -2.60617072, -3.54039369, -2.81142371,
       -2.38154765, -3.35176432, -4.22263284, -2.85837621, -3.56297889,
       -2.92803185, -3.56412279, -2.79003898, -2.13774464, -1.74911698,
       -2.49939225, -1.20106592, -0.0675282 ,  1.13711641,  0.69073678,
        1.49858271,  0.30238539,  2.45352897,  1.65260743,  3.02276826,
        5.39840388,  6.21079478,  5.96272438,  4.87368142,  4.91506842,
        5.65563803,  5.57862366,  7.59519287,  7.12200265,  6.88474515,
        7.08583168,  4.95653629,  5.45320003,  6.40543129,  5.28645928,
        4.99484262,  4.56229058,  3.91497686,  3.45308118,  5.73552822,
        5.70828312,  5.71386026,  7.20837514,  6.99787072,  7.46076125,
        8.04659924,  8.26748513,  9.79241354,  9.4379378 ,  7.91762497,
        7.56134555,  7.92131285,  6.79767604,  6.48710415,  7.13343463,
        7.50912407,  7.16539704,  7.17244582,  6.97024967,  5.2509996 ,
        5.30417648,  4.68437961,  5.0177496 ,  4.95994981,  6.09600679,
        5.07968535,  3.96360113,  3.75556763,  2.52363477,  2.39499007,
        4.02080414,  2.27589845,  2.99881177,  2.36529629,  2.78644801,
        1.80106143,  1.75763401,  2.29626043,  2.86074076,  3.02107294,
        3.18710387,  2.80867405,  3.94195381,  4.16396913,  2.28553528,
        3.79434026,  5.01556549,  6.00965409,  7.18310335,  6.68105476,
        6.17725351,  5.66948074,  6.5077223 ,  6.97064134,  6.97287429,
        8.03404648,  7.52526574,  7.17763844,  5.82849234,  6.03117604,
        5.69488505,  3.58141487,  4.85340842,  5.34894712,  4.59532311,
        6.10019408,  6.00352532,  5.73546038,  6.97542338,  7.0912687 ,
        6.9798788 ,  8.4305599 ,  9.8459736 ,  9.60978579, 10.57637461,
       10.73668396,  9.69485929, 10.38395516, 11.36336338, 12.17683367,
       14.26241058, 12.13773145, 13.00059902, 13.20234373, 13.74922893,
       15.83810596, 13.94017774, 13.60306366, 13.90716223, 14.63460882,
       14.3919169 , 13.52539842, 12.97805865, 13.14501304, 12.28211262,
       11.1807492 , 11.6407465 , 11.58184839, 11.88524257, 11.91219437,
       11.34360673, 11.41657134, 14.1683999 , 14.81621407, 14.15467697,
       13.55231918, 13.76754797, 12.80212146, 12.82578768, 11.82489955,
       11.57263271, 10.96631467,  9.90016628, 10.70765092,  9.74392442,
       10.27627626,  9.8150736 , 10.85924213, 11.13317111, 11.12161069,
       12.24068524, 10.5835783 , 10.70067539, 11.37105768, 10.6343708 ,
       12.56810023, 11.497757  ,  9.95350276,  9.39774844, 10.53234193,
        8.51061752,  7.94512862,  7.3448837 ,  5.73231441,  3.16054665,
        3.58260687,  4.0697326 ,  4.41462296,  5.30135361,  4.63142467,
        3.99147635,  5.85880958,  6.34297235,  6.81472854,  5.63493745,
        6.16140642,  7.33179497,  8.55253337,  7.19346224,  6.39337871,
        5.4748459 ,  4.84613659,  3.72755146,  3.34835593,  3.44181746,
        3.88551984,  5.2802603 ,  6.34270359,  8.64111537,  9.46141703,
       10.01054801,  9.57430283,  9.91367504, 11.25878249, 11.96680951,
       12.50168792, 13.75768347, 14.16221838, 14.3263582 , 14.31562868,
       15.19374288, 14.7600818 , 15.3760113 , 14.97896492, 17.1607779 ,
       17.878575  , 16.63040137, 17.0136007 , 17.3882358 , 15.98494235,
       15.58329047, 13.98019013, 13.50297273, 13.74632653, 12.95940078,
       13.79227236, 13.2810606 , 12.74194887, 14.09310746, 14.94963949,
       15.99535383, 15.14370129, 15.49349416, 16.35879882, 15.14646264,
       16.26710306, 15.24774197, 16.26142951, 16.08461415, 17.17449352,
       17.14627181, 16.27949238, 15.73370298, 15.59387885, 16.88516062,
       15.77413571, 15.4977976 , 14.10955263, 13.0086392 , 12.58363645,
       13.52172609, 13.10290919, 13.34613824, 12.77612736, 13.34593588,
       14.04648451, 14.81864657, 15.10879351, 17.2413906 , 18.55385081,
       18.52524038, 17.48524376, 17.58863619, 17.01259317, 17.27572279,
       17.12693745, 16.72030729, 16.66448996, 16.53693136, 15.30945173,
       14.28257757, 13.92546845, 14.34310097, 15.04712182, 13.71410016,
       12.14414038, 10.41610951, 10.32870421, 11.1136082 , 10.72000385,
       12.45705826, 13.91806901, 15.06015516, 16.71476865, 16.46215077,
       17.64314984, 18.3853737 , 18.46557593, 18.21512902, 16.78894611,
       19.14136502, 19.72682617, 19.78399777, 18.5530771 , 19.18598746,
       19.08978484, 18.44821176, 19.602576  , 19.6409283 , 19.46887729,
       20.16050805, 20.98574552, 22.47905634, 23.15344483, 23.92183549,
       24.23149451, 24.11473377, 22.82201664, 24.13628181, 24.59253408,
       27.25299229, 27.24803585, 27.64377915, 27.17924975, 29.08627341],      dtype=float64)

round

m.dist.normal(0,1, name = 'a', shape=(100,4, 5), sample = True, to_jax=False).round()

SampledData([[[ 2.,  1., -0., -0., -0.],
        [-2., -1., -1., -1., -1.],
        [-1.,  1., -0.,  2., -0.],
        [ 1., -0.,  1.,  1.,  1.]],

       [[-0., -2.,  3.,  0.,  1.],
        [-0., -2., -0.,  1.,  2.],
        [ 2.,  1.,  1.,  1., -2.],
        [-1., -1., -0.,  0.,  1.]],

       [[ 1., -0., -0., -1., -0.],
        [-2.,  1.,  0.,  2., -2.],
        [-2.,  0., -0., -1., -1.],
        [-1.,  1., -0., -1.,  1.]],

       ...,

       [[ 1.,  1., -2.,  1.,  1.],
        [-0.,  2.,  1., -0., -1.],
        [ 0., -1.,  1., -0., -1.],
        [ 0., -1.,  1.,  1., -1.]],

       [[-1., -0., -1., -0.,  0.],
        [ 1., -1., -1.,  0.,  1.],
        [-2.,  1.,  0.,  0., -1.],
        [ 1., -1., -1., -1.,  1.]],

       [[ 2.,  2., -0., -0.,  2.],
        [-1., -1.,  1.,  1.,  0.],
        [-2.,  2., -0., -2., -1.],
        [ 2., -1.,  0., -1., -1.]]], dtype=float64)

clip

m.dist.normal(0,1, name = 'a', shape=(100,4), sample = True, to_jax=False).clip(0, 1)

SampledData([[0.33758575, 0.81261886, 0.        , 0.0561531 ],
       [0.        , 0.50405656, 1.        , 1.        ],
       [1.        , 0.6289524 , 0.        , 0.        ],
       [1.        , 0.        , 0.        , 0.17009709],
       [0.01349144, 0.        , 0.        , 1.        ],
       [0.06701516, 0.50894853, 0.19140889, 0.47035825],
       [0.        , 0.05196079, 0.        , 1.        ],
       [0.        , 1.        , 1.        , 0.        ],
       [0.        , 0.31031892, 0.        , 0.63858941],
       [0.        , 0.        , 0.        , 0.        ],
       [0.        , 1.        , 0.        , 0.99521563],
       [0.71407616, 0.01258172, 0.91458146, 0.82607203],
       [0.        , 0.        , 0.69936457, 0.48987761],
       [0.40045123, 0.        , 0.76719189, 1.        ],
       [1.        , 0.        , 0.17170408, 1.        ],
       [1.        , 0.66954711, 0.28250865, 0.        ],
       [0.64829073, 1.        , 0.        , 1.        ],
       [0.        , 0.        , 0.85950908, 0.6643904 ],
       [0.        , 0.78656933, 0.        , 0.02980835],
       [0.        , 0.02465924, 0.62679637, 0.1356336 ],
       [0.02851979, 0.        , 0.        , 0.        ],
       [0.12362045, 0.        , 0.59609672, 0.43117556],
       [0.57591699, 0.        , 0.8953026 , 0.20035001],
       [0.53105593, 1.        , 0.12410884, 0.        ],
       [0.        , 0.44561766, 0.20289503, 0.        ],
       [0.        , 0.        , 1.        , 1.        ],
       [0.        , 0.        , 0.        , 0.        ],
       [0.        , 0.        , 0.        , 1.        ],
       [1.        , 0.        , 0.52086806, 1.        ],
       [1.        , 0.        , 1.        , 0.36805038],
       [0.18537994, 0.        , 1.        , 0.83624705],
       [0.        , 0.        , 0.89333711, 1.        ],
       [0.        , 0.        , 0.        , 0.        ],
       [1.        , 0.        , 0.36580324, 0.        ],
       [1.        , 0.69137571, 0.        , 0.        ],
       [0.        , 1.        , 1.        , 0.63215534],
       [0.        , 0.        , 0.        , 0.        ],
       [0.        , 0.31220634, 0.        , 0.        ],
       [0.55329676, 0.14350319, 0.2602471 , 0.        ],
       [0.        , 0.        , 0.        , 0.06874563],
       [1.        , 0.        , 0.30758299, 0.67382679],
       [0.91205028, 0.        , 1.        , 0.        ],
       [1.        , 0.        , 0.        , 0.        ],
       [0.        , 0.2125325 , 1.        , 0.74890636],
       [1.        , 0.        , 0.55070669, 1.        ],
       [0.38876432, 0.        , 0.45821444, 0.01173775],
       [0.        , 0.        , 0.52838753, 0.        ],
       [0.02299939, 0.        , 0.        , 0.        ],
       [0.        , 0.        , 0.98493604, 0.        ],
       [0.        , 0.40157583, 0.        , 1.        ],
       [0.        , 1.        , 0.73075902, 0.        ],
       [0.51156743, 0.05923208, 0.99182856, 1.        ],
       [1.        , 0.67693324, 1.        , 0.        ],
       [0.        , 0.18563779, 0.        , 0.        ],
       [0.        , 0.        , 0.        , 0.75495684],
       [0.70398882, 1.        , 0.38379665, 0.47907049],
       [1.        , 0.        , 0.80100268, 0.        ],
       [0.        , 0.        , 0.        , 0.        ],
       [0.        , 0.33956669, 0.        , 0.35875572],
       [0.52271633, 0.12862384, 0.23301847, 0.50516976],
       [0.        , 0.        , 0.        , 0.        ],
       [0.08423882, 0.        , 0.        , 1.        ],
       [0.        , 1.        , 0.        , 0.        ],
       [0.        , 0.        , 0.42210017, 0.25908631],
       [0.        , 0.29107612, 1.        , 0.        ],
       [0.17627117, 0.38822366, 0.04841842, 0.06205096],
       [0.        , 0.40577576, 0.28609854, 0.        ],
       [0.22514519, 0.        , 0.52958913, 0.        ],
       [0.67239768, 0.12100293, 0.23194665, 0.        ],
       [0.        , 0.        , 0.        , 0.        ],
       [0.        , 1.        , 0.47352708, 1.        ],
       [0.60645107, 0.        , 1.        , 0.        ],
       [0.        , 0.59613553, 0.        , 0.0914707 ],
       [0.        , 0.        , 0.51968361, 0.87649035],
       [1.        , 0.        , 0.18200028, 0.40290424],
       [0.        , 0.        , 1.        , 0.        ],
       [0.85825188, 0.43104956, 1.        , 0.        ],
       [0.        , 0.        , 1.        , 0.89211317],
       [0.        , 0.        , 0.        , 0.        ],
       [0.95091847, 0.        , 0.        , 0.34159101],
       [0.        , 0.        , 0.        , 0.        ],
       [1.        , 0.        , 0.54802992, 0.        ],
       [0.        , 0.        , 0.        , 0.        ],
       [0.03549006, 0.82838329, 0.18770956, 0.        ],
       [0.75321014, 1.        , 0.        , 0.        ],
       [0.14539788, 0.        , 0.44861854, 0.        ],
       [0.        , 0.19613261, 0.3427868 , 0.2322019 ],
       [1.        , 0.        , 0.        , 0.        ],
       [0.79410998, 0.        , 0.50718649, 0.41226117],
       [1.        , 0.        , 0.        , 0.        ],
       [0.72832797, 0.        , 0.        , 0.27604445],
       [0.98642783, 0.81431391, 0.        , 0.        ],
       [0.        , 0.        , 0.        , 0.04207799],
       [0.        , 0.        , 0.        , 0.56269332],
       [0.        , 0.55403328, 0.        , 1.        ],
       [0.        , 0.99840991, 0.        , 0.        ],
       [1.        , 0.        , 0.        , 0.        ],
       [0.        , 1.        , 1.        , 0.        ],
       [0.18855704, 1.        , 0.        , 0.        ],
       [0.        , 0.        , 0.        , 0.        ]], dtype=float64)

ptp

m.dist.normal(0,1, name = 'a', shape=(100,4), sample = True, to_jax=False).ptp()

Array(5.9220151, dtype=float64)

sort

m.dist.normal(0,1, name = 'a', shape=(100,4), sample = True, to_jax=False).sort()

SampledData([[-1.1323815 , -0.63807202, -0.44975116,  0.22471516],
       [-1.27617015, -0.6500647 , -0.12255122,  0.50445291],
       [-1.21007533, -0.32533942,  1.26811939,  2.18810049],
       [-0.74875767, -0.04956313,  0.33312353,  0.70715738],
       [-1.32527497, -0.34795518,  0.47412506,  1.64227041],
       [-0.29837047, -0.1142262 ,  0.76181966,  1.1457828 ],
       [-1.92315667, -0.6169582 ,  0.76364436,  1.03158864],
       [-1.66533048, -0.69561994,  1.03768612,  2.12271236],
       [-1.40753559, -0.90761245, -0.70447013,  0.41437365],
       [-1.08842609,  0.04026392,  0.77416832,  1.93885302],
       [-1.09676347, -0.55351501,  0.86703771,  1.01771839],
       [-0.94279173, -0.86416154, -0.43669835,  2.05323757],
       [-0.52090906,  0.23514424,  0.86162409,  1.61876968],
       [-2.66140293, -0.41552552,  0.06623361,  0.23994697],
       [-1.24513045, -0.18585894, -0.01641363,  0.5755547 ],
       [-0.17589499, -0.08666909,  0.58873057,  1.6400648 ],
       [-0.58649591, -0.5823018 ,  0.073012  ,  0.54291569],
       [-0.71914875,  0.2598152 ,  0.41205703,  1.40477226],
       [-0.57438946, -0.34138415, -0.1777548 , -0.05995879],
       [-0.87395951, -0.14382565, -0.02016103,  1.55613941],
       [-1.43487406,  0.88445526,  1.14595141,  1.73436189],
       [-0.87260751, -0.37362484,  0.57125226,  0.97669436],
       [-0.53207412,  0.14533391,  0.19304936,  1.36028199],
       [-1.77455261, -1.33787283, -0.68186274,  0.2603306 ],
       [-0.68232016, -0.02157708,  0.17184814,  0.75935577],
       [-0.61442019,  0.18946737,  1.35123853,  1.66984942],
       [-1.10780081, -0.8603615 , -0.56136483,  0.39290491],
       [-1.74558731, -0.17025859, -0.16634757,  0.74158501],
       [-0.93146873, -0.8357341 ,  1.51305541,  1.94058035],
       [-1.66809762,  0.21180132,  0.30429616,  1.99965543],
       [-1.98966133, -0.99595703, -0.71128608, -0.41583612],
       [-1.15277503, -0.34833257, -0.17540759, -0.14308282],
       [-1.2299368 , -0.64747339,  0.06157168,  0.55021687],
       [-1.10708979, -0.32271928, -0.0618001 ,  0.64719514],
       [-0.74300958,  0.11876789,  0.40040791,  0.629902  ],
       [-1.60013626, -0.50004572,  0.05439235,  1.18169037],
       [-0.78008569,  0.21080881,  0.6143533 ,  1.21915749],
       [-0.54944772,  0.19928205,  0.47043941,  1.23475381],
       [-1.48098382, -0.52732174,  0.11206507,  0.26536663],
       [-0.18174199,  0.13945679,  0.38172068,  1.12012655],
       [-1.43882864, -0.13988237,  0.4616507 ,  0.99855577],
       [-0.96294223, -0.85694397, -0.44443077,  0.83429707],
       [-0.45966903, -0.44027453,  0.64988196,  1.05965297],
       [-0.25548378,  0.36662383,  1.30810486,  1.33924914],
       [-0.73231657, -0.11402064,  0.47443387,  0.632928  ],
       [-0.09081861,  0.4431942 ,  0.65144286,  1.34171873],
       [-1.81178476, -0.08900717,  0.15837666,  0.20502665],
       [-1.17041247,  0.79878975,  1.60747929,  1.91052768],
       [-1.19160338, -0.02536563,  0.13848709,  0.92664348],
       [-0.80875579, -0.26237871,  0.67731894,  0.92152762],
       [-0.08961022,  0.07687657,  0.25909043,  0.95654685],
       [-1.96145792, -1.1294143 , -0.71452188,  0.56823666],
       [-1.5583682 , -1.08401915,  0.49103637,  0.7091078 ],
       [-0.21542227,  0.1442121 ,  0.36409956,  0.88886572],
       [-0.97284523, -0.75370989, -0.59717967, -0.49514576],
       [-0.85066418, -0.71872636, -0.02617978,  0.05900249],
       [-1.44606038, -0.59558755, -0.24618777,  1.65432316],
       [-1.87885232, -1.02814703, -0.74329775, -0.47885007],
       [-1.12288137, -0.373728  ,  0.36250207,  0.98093018],
       [-0.74766809,  0.18678436,  0.58537222,  1.2305321 ],
       [-1.93279537, -0.98028744, -0.10585916,  1.08157545],
       [-0.02611873,  0.23145597,  0.33234918,  1.40556336],
       [-1.08891565,  0.75880968,  0.80008494,  1.29037918],
       [-2.79346275, -0.39729164,  0.44300361,  0.66119545],
       [-0.10191365, -0.09423111,  1.28089744,  1.55359112],
       [-1.56318536, -0.70325574,  0.05856932,  1.61863463],
       [-0.56253344, -0.08694989,  0.06012768,  1.32360441],
       [-1.92344634, -1.28173607,  0.61302603,  1.54713491],
       [-0.54149942,  0.05379851,  0.24442708,  0.52000964],
       [-0.61840711,  0.4999324 ,  2.00876219,  2.10882938],
       [-1.65537648, -1.57336167,  0.2848829 ,  1.10225769],
       [-0.4534426 , -0.42237189, -0.12357035,  0.18310777],
       [-1.01657921, -0.1168166 ,  0.31306684,  0.64914341],
       [-1.92393566, -1.74753132,  0.2534945 ,  1.43536324],
       [-0.26632396,  0.04946166,  0.21869778,  0.65609238],
       [ 0.35809737,  0.552411  ,  0.55419579,  1.49045306],
       [-2.05097437, -1.81060917, -0.0652534 ,  0.12081531],
       [-1.29924657, -0.7708373 ,  0.63447425,  1.20215794],
       [-0.32078247, -0.08686784,  1.67004098,  1.85671816],
       [-1.15529814, -0.04070542,  1.57176114,  2.70039694],
       [-1.14133176, -0.32746484,  0.49448502,  1.84300859],
       [-0.47157391, -0.42620746, -0.06690037,  0.67451914],
       [ 0.24057541,  0.93696202,  1.33635603,  1.81280427],
       [-0.6505614 , -0.05282716,  0.65093691,  0.85385701],
       [-1.39013198, -1.00928143, -0.01436457,  1.39999589],
       [-1.1228363 , -0.29303232,  0.2334246 ,  2.04846856],
       [-0.45475748,  0.20577837,  0.7481044 ,  1.38593181],
       [-0.57302822, -0.54105272, -0.45587357, -0.13690456],
       [-0.0297334 ,  0.42686575,  0.69887721,  0.78391598],
       [-0.3840571 ,  0.68215728,  0.70896884,  2.27932682],
       [-0.67610822, -0.5191818 ,  0.46778871,  0.97189902],
       [-0.27924575,  0.41148137,  1.19697043,  1.32122535],
       [-0.33151752,  0.49015256,  0.96313276,  1.11437563],
       [-2.09146766, -1.36030081, -0.43406441, -0.33255515],
       [-0.26903957,  0.00393225,  0.65052782,  1.16668459],
       [-1.41165093, -1.22824515,  0.72229508,  0.8396216 ],
       [-1.21996398, -0.82704246, -0.25618686,  1.18649538],
       [-2.61100253, -2.38347709, -0.03393367,  1.01318012],
       [-0.61988964, -0.408882  ,  1.02100823,  2.19177876],
       [-0.33780032,  0.01236119,  0.25686123,  0.54579431]],      dtype=float64)

.at[]

Since the sampled object wraps a JAX array, you can use JAX’s .at[] syntax for in-place style updates (which return a new object):

s = m.dist.normal(0,1, name = 'a', shape=(100,4, 5), sample = True, to_jax=False)
s_new = s.at[0, 0, 0].set(0.0)

reshape

m.dist.normal(0,1, name = 'a', shape=(100,4, 5), sample = True, to_jax=False).reshape((100, 20))

SampledData([[ 0.3135452 ,  2.20980842, -0.81509743, ..., -0.46802198,
         0.15228599,  0.26250653],
       [-1.8323697 ,  0.57090966, -0.72098376, ..., -1.61941109,
         0.2327595 ,  0.23252358],
       [ 0.80899582, -1.27283438, -0.2193537 , ..., -1.14804418,
        -0.52030859,  1.23155218],
       ...,
       [ 0.57055981,  0.17119572, -0.5081537 , ...,  1.08977935,
         0.40611105,  0.27900972],
       [ 0.68851658,  0.91047305,  0.87893022, ..., -1.0637584 ,
        -0.57934575,  0.23224165],
       [ 0.78668513, -2.03077693, -0.70893224, ..., -0.38449669,
        -0.64263684, -0.05597747]], dtype=float64)

And many other standard NumPy/JAX array methods like .flatten(), .squeeze(), .astype(), etc., are also available directly on the sampled object.

Notes

Note

• density, surface_3d, and ppc_plot are Plotly-only — passing interactive=False has no effect on them.
• The dimensional requirement is enforced: e.g. corr_heatmap, boxplot, violinplot, pairplot, traceplot, scatter3d, and timeseries raise a ValueError if the array is not the expected shape.
• Plotting methods live on the SampledData wrapper only. If you sampled with to_jax=True (or later called .to_jax()), convert back with SampledData(arr) to plot, or re-sample without to_jax.