Finish random variable port to backend

Author: marvinpfoertner

src/probnum/backend/__init__.py

@@ -18,6 +18,7 @@
     "is_floating",
     "finfo",
     # Shape Arithmetic
+    "reshape",
     "atleast_1d",
     "atleast_2d",
     "broadcast_arrays",

src/probnum/backend/_core/__init__.py

@@ -27,6 +27,7 @@
 finfo = _core.finfo
 
 # Shape Arithmetic
+reshape = _core.reshape
 atleast_1d = _core.atleast_1d
 atleast_2d = _core.atleast_2d
 broadcast_arrays = _core.broadcast_arrays

src/probnum/backend/_core/_jax.py

@@ -33,6 +33,7 @@
     ones_like,
     pi,
     promote_types,
+    reshape,
     sin,
     single,
     sqrt,

src/probnum/backend/_core/_numpy.py

@@ -32,6 +32,7 @@
     ones_like,
     pi,
     promote_types,
+    reshape,
     sin,
     single,
     sqrt,

src/probnum/backend/_core/_torch.py

@@ -26,6 +26,7 @@
     maximum,
     pi,
     promote_types,
+    reshape,
     sin,
     sqrt,
 )

src/probnum/randvars/_arithmetic.py

@@ -123,19 +123,21 @@ def _rv_binary_op(rv1: _RandomVariable, rv2: _RandomVariable) -> _RandomVariable
 
 
 def _make_rv_binary_op_result_shape_dtype_sample_fn(op_fn, rv1, rv2):
-    seed = backend.random.seed(1)
-    sample_fn = lambda sample_shape: op_fn(
-        rv1.sample(seed=seed, sample_shape=sample_shape),
-        rv2.sample(seed=seed, sample_shape=sample_shape),
-    )
+    def sample(seed, sample_shape):
+        seed1, seed2, _ = backend.random.split(seed, 3)
+
+        return op_fn(
+            rv1.sample(seed=seed1, sample_shape=sample_shape),
+            rv2.sample(seed=seed2, sample_shape=sample_shape),
+        )
 
     # Infer shape and dtype
-    infer_sample = sample_fn(())
+    infer_sample = sample(backend.random.seed(1), ())
 
     shape = infer_sample.shape
     dtype = infer_sample.dtype
 
-    return shape, dtype, sample_fn
+    return shape, dtype, sample
 
 
 def _generic_rv_add(rv1: _RandomVariable, rv2: _RandomVariable) -> _RandomVariable:

src/probnum/randvars/_random_variable.py

@@ -313,7 +313,10 @@ def var(self) -> ArrayType:
         """
         if self.__var is None:
             try:
-                var = np.diag(self.cov).reshape(self.__shape).copy()
+                var = backend.reshape(
+                    backend.diag(self.cov),
+                    self.__shape,
+                ).copy()
             except NotImplementedError as exc:
                 raise NotImplementedError from exc
         else:
@@ -514,8 +517,12 @@ def quantile(self, p: ArrayType) -> ArrayType:
         return quantile
 
     def __getitem__(self, key: ArrayLikeGetitemArgType) -> "RandomVariable":
+        # Shape inference
+        # For simplicity, this should not be computed using backend, but rather in numpy
+        shape = np.broadcast_to(np.empty(()), self.shape)[key].shape
+
         return RandomVariable(
-            shape=np.empty(shape=self.shape)[key].shape,
+            shape=shape,
             dtype=self.dtype,
             sample=lambda rng, size: self.sample(rng, size)[key],
             mode=lambda: self.mode[key],
@@ -557,8 +564,13 @@ def transpose(self, *axes: int) -> "RandomVariable":
         axes :
             See documentation of :meth:`numpy.ndarray.transpose`.
         """
+
+        # Shape inference
+        # For simplicity, this should not be computed using backend, but rather in numpy
+        shape = np.broadcast_to(np.empty(()), self.shape).transpose(*axes).shape
+
         return RandomVariable(
-            shape=np.empty(shape=self.shape).transpose(*axes).shape,
+            shape=shape,
             dtype=self.dtype,
             sample=lambda rng, size: self.sample(rng, size).transpose(*axes),
             mode=lambda: self.mode.transpose(*axes),