Out of memory error

[mscroggs]

Describe the Bug

The test test_mean_stdev_two_node_graph currently gives an out of memory error when sampling 10^6 values.

To Reproduce

Uncomment the xfail, then run:

python -m pytest tests/test_graph/test_algorithms.py::test_mean_stdev_two_node_graph

Expected Behaviour

Test passes

Actual Behaviour

jaxlib.xla_extension.XlaRuntimeError: FAILED_PRECONDITION: Buffer Definition Event: Error dispatching computation: Error dispatching computation: Error dispatching computation: Error preparing computation: Out of memory allocating 4000000000000 bytes.

[mscroggs]

The test no longer gives a memory error following the fixing of the shapes of sample in #70