Fix bug in output sampler on GPU (#170)

gianlucadetommaso · web-flow · commit ae5c25483dd6 · 2023-12-19T12:48:41.000+01:00
diff --git a/fortuna/prob_model/predictive/base.py b/fortuna/prob_model/predictive/base.py
@@ -563,7 +563,8 @@ def _sample(key, _inputs):
             )
             if distribute:
                 outputs = jnp.stack(
-                    list(map(lambda key: _sample(shard_prng_key(key), inputs), keys))
+                    list(map(lambda key: _sample(shard_prng_key(key), inputs), keys)),
+                    axis=1,
                 )
                 outputs = self._unshard_ensemble_arrays(outputs)
             else:

Original file line number	Diff line number	Diff line change
`@@ -563,7 +563,8 @@ def _sample(key, _inputs):`
`563`	`563`	`)`
`564`	`564`	`if distribute:`
`565`	`565`	`outputs = jnp.stack(`
`566`		`- list(map(lambda key: _sample(shard_prng_key(key), inputs), keys))`
	`566`	`+ list(map(lambda key: _sample(shard_prng_key(key), inputs), keys)),`
	`567`	`+ axis=1,`
`567`	`568`	`)`
`568`	`569`	`outputs = self._unshard_ensemble_arrays(outputs)`
`569`	`570`	`else:`