# Any help appreciated!

**URL:** https://discuss.bayesflow.org/t/any-help-appreciated/237
**Category:** General
**Created:** [April 15, 2026, 7:47am UTC](https://discuss.bayesflow.org/t/any-help-appreciated/237 "2026-04-15T07:47:20Z")
**Posts on this page:** 6
**Page:** 1

<div class="post-metadata">

### Author: ![Will](https://avatars.discourse-cdn.com/v4/letter/w/90ced4/32.png) [@Will](https://discuss.bayesflow.org/u/Will)
#### Post date: [April 15, 2026, 7:47am UTC](https://discuss.bayesflow.org/t/any-help-appreciated/237/1 "2026-04-15T07:47:20Z")

</div>

Hello! I’m in the process of trying to make a nice basic 3 parameter DDM and I hope to be able to nail the basic version so I can branch out in the future. I had some initial success but there were some flaws in the logic that I wanted to smooth out. Attempting to fix this, I’ve gotten into trouble. I’m quite poor at coding (particularly python) so trying to debug it has been quite difficult. I’m very new to this all so any help on anything in general that stands out as a red flag or anything more specific based on this error that I keep getting would be greatly appreciated. Thank you!

---

<div class="post-metadata">

### Author: ![Will](https://avatars.discourse-cdn.com/v4/letter/w/90ced4/32.png) [@Will](https://discuss.bayesflow.org/u/Will)
#### Post date: [April 15, 2026, 7:52am UTC](https://discuss.bayesflow.org/t/any-help-appreciated/237/2 "2026-04-15T07:52:29Z")

</div>

Script 1 (which worked)

===== 1. Setup =====

import os  
os.environ[“KERAS\_BACKEND”]=“jax”  
os.environ[“TF\_CUDNN\_DETERMINISTIC”]=“1”  
os.environ[“XLA\_PYTHON\_CLIENT\_PREALLOCATE”]=“false”

import numpy as np, bayesflow as bf, keras  
import jax, jax.numpy as jnp

print(“=== JAX GPU check ===”)  
print(“JAX version:”, jax. **version** )  
print(“Devices:”, jax.devices())  
gpu=[d for d in jax.devices() if d.platform==“gpu”]  
print(“GPU detected:” if gpu else “⚠ No GPU detected.”, gpu)  
print(“====================\n”)

===== 2. Priors =====

def prior():  
return dict(  
v=float(np.random.uniform(0.2,2.5)),  
a=float(np.random.uniform(0.5,2.5)),  
t0=float(np.random.gamma(2,0.25)),  
)

===== 3. JAX Simulator =====

@jax.jit(static\_argnames=[“n\_obs”,“max\_steps”,“dt”])  
def \_simulate\_batch(key,v,a,t0,n\_obs,max\_steps,dt):  
z=0.5\*a  
time\_steps=jnp.arange(max\_steps,dtype=jnp.float32)\*dt  
key\_noise,key\_tie=jax.random.split(key)  
noise=jax.random.normal(key\_noise,(n\_obs,max\_steps))\*jnp.sqrt(dt)

traj=z+jnp.cumsum(v\*dt+noise,axis=1)  
up,low=traj\>=a,traj\<=0.0

up\_idx,low\_idx=jnp.argmax(up,1),jnp.argmax(low,1)  
hit\_up,hit\_low=jnp.any(up,1),jnp.any(low,1)

tie=(up\_idx==low\_idx)&hit\_up&hit\_low  
rand=jax.random.bernoulli(key\_tie,0.5,shape=up\_idx.shape)  
up\_first=jnp.where(tie,rand,up\_idx\<low\_idx)

up\_rt=jnp.where(hit\_up,time\_steps[up\_idx]+t0,jnp.nan)  
low\_rt=jnp.where(hit\_low,time\_steps[low\_idx]+t0,jnp.nan)

rt=jnp.where(hit\_up&hit\_low,  
jnp.where(up\_first,up\_rt,low\_rt),  
jnp.where(hit\_up,up\_rt,  
jnp.where(hit\_low,low\_rt,jnp.nan)))

resp=jnp.where(hit\_up&hit\_low,  
jnp.where(up\_first,1.0,0.0),  
jnp.where(hit\_up,1.0,  
jnp.where(hit\_low,0.0,-1.0)))  
return rt,resp

def simulator(v,a,t0,n\_obs,dt=0.001,max\_rt=2.0,min\_rt=0.2):  
key=jax.random.PRNGKey(np.random.randint(0,2\*\*31-1))  
rt,resp=\_simulate\_batch(key,v,a,t0,n\_obs,int(max\_rt/dt),dt)

rt,resp=np.asarray(rt),np.asarray(resp)  
valid=(~np.isnan(rt))&(resp\>=0)&(rt\>=min\_rt)&(rt\<=max\_rt)

rt\_v,resp\_v=rt[valid],resp[valid]  
if len(rt\_v)\>0:  
idx=np.random.choice(len(rt\_v),n\_obs,  
replace=len(rt\_v)\<n\_obs)  
rt\_v,resp\_v=rt\_v[idx],resp\_v[idx]  
else:  
rt\_v=np.full(n\_obs,max\_rt,np.float32)  
resp\_v=np.zeros(n\_obs,np.float32)

return dict(x=np.stack([rt\_v,resp\_v],1).astype(np.float32))

===== 4. Simulator Object =====

def meta():  
return dict(n\_obs=np.int32(np.random.randint(120,161)))

simulator\_obj=bf.make\_simulator([prior,simulator],meta\_fn=meta)

===== 5. Workflow =====

adapter=(bf.Adapter()  
.constrain(“a”,lower=0)  
.constrain(“t0”,lower=0)  
.as\_set(“x”)  
.concatenate([“v”,“a”,“t0”],into=“inference\_variables”)  
.rename(“x”,“summary\_variables”))

workflow=bf.BasicWorkflow(  
simulator=simulator\_obj,  
adapter=adapter,  
inference\_network=bf.networks.FlowMatching(),  
summary\_network=bf.networks.SetTransformer(summary\_dim=256),  
inference\_variables=[“v”,“a”,“t0”],  
summary\_variables=[“x”],  
)

===== 6. Training =====

opt=keras.optimizers.AdamW(learning\_rate=1e-4)  
history=workflow.fit\_online(  
epochs=200,  
num\_batches\_per\_epoch=2500,  
batch\_size=32,  
optimizer=opt,  
)

===== 7. Save =====

workflow.approximator.save(“accuracy\_ddm\_9\_approximator.keras”)  
print(“Model saved.”)

---

<div class="post-metadata">

### Author: ![Will](https://avatars.discourse-cdn.com/v4/letter/w/90ced4/32.png) [@Will](https://discuss.bayesflow.org/u/Will)
#### Post date: [April 15, 2026, 7:56am UTC](https://discuss.bayesflow.org/t/any-help-appreciated/237/3 "2026-04-15T07:56:02Z")

</div>

Script 2: Has not been working, tried to make it more principled but I am having issues with it. (All I changed when pasting here was some of the formatting)

---- Section 1 ----- Import/Set Up  
import os  
os.environ[“KERAS\_BACKEND”] = “jax”  
os.environ[“TF\_CUDNN\_DETERMINISTIC”] = “1”  
os.environ[“XLA\_PYTHON\_CLIENT\_PREALLOCATE”] = “false”

import numpy as np  
import bayesflow as bf  
import keras  
from scipy.stats import truncnorm

import jax  
import jax.numpy as jnp

print(“=== JAX GPU check ===”)  
print(“JAX version:”, jax. **version** )  
print(“Available devices:”, jax.devices())

gpu\_devices = [d for d in jax.devices() if d.platform == ‘gpu’]  
if gpu\_devices:  
print(f"GPU detected: {gpu\_devices}“)  
else:  
print(”⚠ No GPU detected. JAX will run on CPU.“)  
print(”====================\n")

print(“JAX devices:”, jax.devices())

---- Section 2 ---- Priors

def truncated\_normal(mean, sd, lower, upper):  
a = (lower - mean) / sd  
b = (upper - mean) / sd  
return truncnorm.rvs(a, b, loc=mean, scale=sd)

def prior():  
return dict(  
v=np.float32(truncated\_normal(mean=1.75, sd=0.70, lower=0.0, upper=2.5)),  
a=np.float32(truncated\_normal(mean=1.75, sd=0.50, lower=0.5, upper=2.5)),  
t0=np.float32(truncated\_normal(mean=0.60, sd=0.15, lower=0.2, upper=1.0)),  
)

---- Section 3 ---- Simulator

@jax.jit(static\_argnames=[“n\_sim”, “max\_steps”, “dt”])  
def \_simulate\_batch(key, v, a, t0, n\_sim, max\_steps, dt):  
z = 0.5 \* a  
time\_steps = jnp.arange(max\_steps, dtype=jnp.float32) \* dt

```
key_noise, key_tie = jax.random.split(key)
s = 1.0
noise = jax.random.normal(key_noise, (n_sim, max_steps)) * (s * jnp.sqrt(dt))

increments = v * dt + noise
trajectories = z + jnp.cumsum(increments, axis=1)

upper_cross = trajectories >= a
lower_cross = trajectories <= 0.0

upper_idx = jnp.argmax(upper_cross, axis=1)
lower_idx = jnp.argmax(lower_cross, axis=1)

hit_upper = jnp.any(upper_cross, axis=1)
hit_lower = jnp.any(lower_cross, axis=1)

tie = (upper_idx == lower_idx) & hit_upper & hit_lower
rand = jax.random.bernoulli(key_tie, 0.5, shape=upper_idx.shape)
upper_first = jnp.where(tie, rand, upper_idx < lower_idx)

upper_rt = jnp.where(hit_upper, time_steps[upper_idx] + t0, jnp.nan)
lower_rt = jnp.where(hit_lower, time_steps[lower_idx] + t0, jnp.nan)

rt = jnp.where(
    hit_upper & hit_lower,
    jnp.where(upper_first, upper_rt, lower_rt),
    jnp.where(hit_upper, upper_rt,
              jnp.where(hit_lower, lower_rt, jnp.nan))
)

resp = jnp.where(
    hit_upper & hit_lower,
    jnp.where(upper_first, 1.0, 0.0),
    jnp.where(hit_upper, 1.0,
              jnp.where(hit_lower, 0.0, -1.0))
)

return rt, resp

```

def simulator(v, a, t0, n\_obs, dt=0.001, max\_rt=10.0):  
n\_sim = 300  
max\_steps = int(max\_rt / dt)

```
key = jax.random.PRNGKey(np.random.randint(0, 2**31 - 1))
rt, resp = _simulate_batch(key, v, a, t0, n_sim, max_steps, dt)

rt = np.asarray(rt)
resp = np.asarray(resp)

return dict(rt=rt, resp=resp)

```

---- Section 4 ---- Preprocessor

def preprocess\_simulator\_output(sim\_output, n\_obs):  
rt = sim\_output[“rt”].reshape(-1)  
resp = sim\_output[“resp”].reshape(-1)

```
valid = (
    (~np.isnan(rt)) &
    (resp >= 0) &
    (rt >= 0.2) &
    (rt <= 2.0)
)

rt_valid = rt[valid]
resp_valid = resp[valid]

if len(rt_valid) == 0:
    mask = ~np.isnan(rt)
    rt_valid = rt[mask]
    resp_valid = resp[mask]

idx = np.random.choice(len(rt_valid), size=n_obs, replace=True)
rt_sampled = rt_valid[idx]
resp_sampled = resp_valid[idx]

log_rt = np.log(rt_sampled)

x = np.stack([log_rt, resp_sampled], axis=1).astype(np.float32)

return dict(x=x)

```

# ---- Section 5 ---- Meta + Full Simulator

def meta():  
return dict(n\_obs=np.int32(160))

def full\_simulator():  
theta = prior()  
n\_obs = int(meta()[“n\_obs”])

```
sim_raw = simulator(theta["v"], theta["a"], theta["t0"], n_obs)
x_dict = preprocess_simulator_output(sim_raw, n_obs)

return dict(
    v=theta["v"],
    a=theta["a"],
    t0=theta["t0"],
    x=x_dict["x"],
)

```

---- Section 6 ---- Workflow (FIXED SHAPES + NAMES)

adapter = (  
bf.Adapter()  
.constrain(“v”, lower=0.0, upper=2.5)  
.constrain(“a”, lower=0.5, upper=2.5)  
.constrain(“t0”, lower=0.2, upper=1.0)  
.as\_set(“x”) # treat x as set  
.concatenate([“v”, “a”, “t0”], into=“inference\_variables”)  
.rename(“x”, “summary\_variables”) # ← critical rename  
)

workflow = bf.BasicWorkflow(  
simulator=simulator\_obj,  
adapter=adapter,  
inference\_network=bf.networks.FlowMatching(),  
summary\_network=bf.networks.SetTransformer(summary\_dim=64),  
inference\_variables=[“v”, “a”, “t0”], # parameter names  
summary\_variables=[“summary\_variables”], # ← matches rename  
)

---- Section 7 ---- Training

optimizer = keras.optimizers.AdamW(learning\_rate=1e-4)

history = workflow.fit\_online(  
epochs=200,  
num\_batches\_per\_epoch=1000,  
batch\_size=32,  
optimizer=optimizer,  
)

---- Section 8 ---- Saving

workflow.approximator.save(“accuracy\_ddm\_10\_approximator.keras”)  
print(“Model saved.”)

---

<div class="post-metadata">

### Author: ![Will](https://avatars.discourse-cdn.com/v4/letter/w/90ced4/32.png) [@Will](https://discuss.bayesflow.org/u/Will)
#### Post date: [April 15, 2026, 8:02am UTC](https://discuss.bayesflow.org/t/any-help-appreciated/237/4 "2026-04-15T08:02:53Z")

</div>

The key traceback is:

ValueError: Incompatible shapes for broadcasting: (96,) and requested shape (1,)

This happens inside:

time = keras.ops.broadcast\_to(time, keras.ops.shape(xz)[:-1] + (1,))

So it looks like time has shape (96,) but is being broadcast to something ending in (1,), which fails.

Again, I’m quite inexperienced so if anything seems like it needs a total redo please let me know, I’ve added the constraints to match the task I want to apply the approximator to. Feel free to reach out to me. Sorry for the ugly posting!

---

<div class="post-metadata">

### Author: ![KLDivergence](https://yyz1.discourse-cdn.com/flex007/user_avatar/discuss.bayesflow.org/kldivergence/32/15_2.png) [@KLDivergence](https://discuss.bayesflow.org/u/KLDivergence)
#### Post date: [April 15, 2026, 12:15pm UTC](https://discuss.bayesflow.org/t/any-help-appreciated/237/5 "2026-04-15T12:15:36Z")

</div>

Hi Will, and welcome to our forums!

First off, since there is no GPU, I suggest instructing the LLM to code up your simulator in pure `numpy`. Also, there is no need to use `as_set` in the adapter, as this is only used for arrays of shape `(num_repeats, )` that are actually **unordered sets** and need to be represented as arrays of shape `(num_repeats, 1)`.

For your setup, you need to make sure that what goes into the networks is:

- `inference_variables` → array of shape `(num_sims, num_params)`
- `summary_variables` → array of shape `(num_sims, num_obs, num_columns)`

I couple of practical suggestions:

1. Always check the shapes of all adapted outputs as a sanity check:

```auto
test_sims = simulator(2)
adapted_sims = adapter(test_sims)
for k, v in adapted_sims.items():
    print(f"Shape of {k} is {v.shape}")

```

1. When you are starting out, you want to train as fast as possible, so don’t use varying `n_obs` and if possible, instruct the LLM to compute summary statistics of the data (e.g., RT quantiles and accuracies per condition). That way, you can forego using a summary network and can pass the “hand-crafted” summaries directly as `inference_conditions` (will reduce your training time to \< 1 minute).
2. Related: do not use online training (`fit_online`) at first, as the total time for any amortized workflow is always:

T\_{total} = T\_{sim} + T\_{train} + T\_{infer}

Online training entangles T\_{train} and T\_{train} and should only be used as a last step in a workflow (i.e., before writing down results for a paper and your model is final). A good heuristic for DDM-like models is that around 10,000 offline simulations are more than sufficient to give you an accurate idea of the recoverability of your model (i.e., using `fit_offline`).

Feel free to check out some of our resources out there on cognitive modeling. For example, this notebooks demonstrates a few workflows with a DDM using pre-built simulators from the excellent `HSSM` library: [bayesflow\_workshops/carney\_comp\_2025/notebooks/ssms\_bayesflow.ipynb at main · bayesflow-org/bayesflow\_workshops · GitHub](https://github.com/bayesflow-org/bayesflow_workshops/blob/main/carney_comp_2025/notebooks/ssms_bayesflow.ipynb)

Let us know if you encounter more roadblocks along the way and happy amortizing!

---

<div class="post-metadata">

### Author: ![Will](https://avatars.discourse-cdn.com/v4/letter/w/90ced4/32.png) [@Will](https://discuss.bayesflow.org/u/Will)
#### Post date: [April 15, 2026, 9:57pm UTC](https://discuss.bayesflow.org/t/any-help-appreciated/237/6 "2026-04-15T21:57:50Z")

</div>

Thanks so much I really appreciate it!