From 7cb316d5bac172999c03f43a6c1e4e8f5d5b3d7d Mon Sep 17 00:00:00 2001
From: holl- <philipp@mholl.de>
Date: Mon, 14 Nov 2022 14:30:29 +0100
Subject: [PATCH 01/29] [math] Improved minimize() error message

---
 phi/math/_functional.py | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/phi/math/_functional.py b/phi/math/_functional.py
index 638832d29..710e38f43 100644
--- a/phi/math/_functional.py
+++ b/phi/math/_functional.py
@@ -1522,9 +1522,12 @@ def native_function(x_flat):
         else:
             y = f(x)
         _, y_tensors = disassemble_tree(y)
-        assert not non_batch(
-            y_tensors[0]), f"Failed to minimize '{f.__name__}' because it returned a non-scalar output {shape(y_tensors[0])}. Reduce all non-batch dimensions, e.g. using math.l2_loss()"
-        return y_tensors[0].sum, (reshaped_native(y_tensors[0], [batch_dims]),)
+        assert not non_batch(y_tensors[0]), f"Failed to minimize '{f.__name__}' because it returned a non-scalar output {shape(y_tensors[0])}. Reduce all non-batch dimensions, e.g. using math.l2_loss()"
+        try:
+            loss_native = reshaped_native(y_tensors[0], [batch_dims])
+        except AssertionError:
+            raise AssertionError(f"Failed to minimize '{f.__name__}' because its output loss {shape(y_tensors[0])} has more batch dimensions than the initial guess {batch_dims}.")
+        return y_tensors[0].sum, (loss_native,)
 
     atol = backend.to_float(reshaped_native(solve.absolute_tolerance, [batch_dims], force_expand=True))
     maxi = backend.to_int32(reshaped_native(solve.max_iterations, [batch_dims], force_expand=True))

From eb18b1cd7658f9b712d36ec69f2d508406cb86bf Mon Sep 17 00:00:00 2001
From: holl- <philipp@mholl.de>
Date: Mon, 14 Nov 2022 14:49:30 +0100
Subject: [PATCH 02/29] [math] Fix print_grad()

---
 phi/math/_functional.py | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/phi/math/_functional.py b/phi/math/_functional.py
index 710e38f43..3174233b7 100644
--- a/phi/math/_functional.py
+++ b/phi/math/_functional.py
@@ -1769,13 +1769,14 @@ def f(x):
         `identity(value)` which when differentiated, prints the gradient vector.
     """
 
-    def print_grad(_x, _y, dx):
-        if all_available(_x, dx):
+    def print_grad(params: dict, _y, dx):
+        param_name, x = next(iter(params.items()))
+        if all_available(x, dx):
             if detailed:
                 print_(dx, name=name)
             else:
                 print(f"{name}:  \t{dx}")
-        return dx,
+        return {param_name: dx}
 
     identity = custom_gradient(lambda x: x, print_grad)
     return identity(value)

From cbbbc27f13e8906fd21574d73186010343d3b9e8 Mon Sep 17 00:00:00 2001
From: holl- <philipp@mholl.de>
Date: Mon, 14 Nov 2022 17:00:52 +0100
Subject: [PATCH 03/29] [math] Add failure message to print_gradient()

---
 phi/math/_functional.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/phi/math/_functional.py b/phi/math/_functional.py
index 3174233b7..a66daf6ae 100644
--- a/phi/math/_functional.py
+++ b/phi/math/_functional.py
@@ -1776,6 +1776,8 @@ def print_grad(params: dict, _y, dx):
                 print_(dx, name=name)
             else:
                 print(f"{name}:  \t{dx}")
+        else:
+            print(f"Cannot print gradient for {param_name}, data not available.")
         return {param_name: dx}
 
     identity = custom_gradient(lambda x: x, print_grad)

From d70502f726275c85b7121d2906b0a2569490db85 Mon Sep 17 00:00:00 2001
From: Philipp Holl <philipp.holl@tum.de>
Date: Mon, 14 Nov 2022 18:44:12 +0100
Subject: [PATCH 04/29] [math] Fix PyTorch einsum() for bool, int

---
 demos/pipe.py                | 3 +--
 phi/math/backend/_backend.py | 4 +++-
 phi/torch/_torch_backend.py  | 2 +-
 3 files changed, 5 insertions(+), 4 deletions(-)

diff --git a/demos/pipe.py b/demos/pipe.py
index 0d25b435d..17bcec56f 100644
--- a/demos/pipe.py
+++ b/demos/pipe.py
@@ -4,8 +4,7 @@
 from phi.flow import *
 
 DT = 1.
-INFLOW_BC = extrapolation.combine_by_direction(normal=1, tangential=0)
-velocity = StaggeredGrid(0, extrapolation.combine_sides(x=(INFLOW_BC, extrapolation.BOUNDARY), y=0), x=50, y=32)
+velocity = StaggeredGrid(0, extrapolation.combine_sides(x=(vec(x=1, y=0), extrapolation.BOUNDARY), y=0), x=50, y=32)
 pressure = None
 
 for _ in view('velocity, pressure', namespace=globals()).range():
diff --git a/phi/math/backend/_backend.py b/phi/math/backend/_backend.py
index 7546a3098..823c498a4 100644
--- a/phi/math/backend/_backend.py
+++ b/phi/math/backend/_backend.py
@@ -134,7 +134,7 @@ def complex_type(self) -> DType:
     def combine_types(self, *dtypes: DType) -> DType:
         return combine_types(*dtypes, fp_precision=self.precision)
 
-    def auto_cast(self, *tensors, bool_to_int=False) -> list:
+    def auto_cast(self, *tensors, bool_to_int=False, int_to_float=False) -> list:
         """
         Determins the appropriate values type resulting from operations involving the tensors as input.
         
@@ -152,6 +152,8 @@ def auto_cast(self, *tensors, bool_to_int=False) -> list:
         result_type = self.combine_types(*dtypes)
         if result_type.kind == bool and bool_to_int:
             result_type = DType(int, 32)
+        if result_type.kind == int and int_to_float:
+            result_type = DType(float, self.precision)
         if result_type.kind in (int, float, complex, bool):  # do not cast everything to string!
             tensors = [self.cast(t, result_type) for t in tensors]
         return tensors
diff --git a/phi/torch/_torch_backend.py b/phi/torch/_torch_backend.py
index d5f6e2dd9..de9148cd0 100644
--- a/phi/torch/_torch_backend.py
+++ b/phi/torch/_torch_backend.py
@@ -147,7 +147,7 @@ def multi_slice(self, tensor, slices: tuple):
     seed = staticmethod(torch.manual_seed)
 
     def einsum(self, equation, *tensors):
-        tensors = self.auto_cast(*tensors)
+        tensors = self.auto_cast(*tensors, bool_to_int=True, int_to_float=True)
         return torch.einsum(equation, *tensors)
 
     def jit_compile(self, f: Callable) -> Callable:

From 00e174e5c78c2afae8b992d676f8ec664ee6514f Mon Sep 17 00:00:00 2001
From: Philipp Holl <philipp.holl@tum.de>
Date: Tue, 15 Nov 2022 11:18:47 +0100
Subject: [PATCH 05/29] [math] Fix dim order in dot()

---
 phi/math/_ops.py | 12 +++++++-----
 1 file changed, 7 insertions(+), 5 deletions(-)

diff --git a/phi/math/_ops.py b/phi/math/_ops.py
index 90013439d..07bde820b 100644
--- a/phi/math/_ops.py
+++ b/phi/math/_ops.py
@@ -1407,10 +1407,11 @@ def dot(x: Tensor,
         result_native = backend.tensordot(x_native, x.shape.indices(x_dims), y_native, y.shape.indices(y_dims))
         result_shape = concat_shapes(remaining_shape_x, remaining_shape_y)
     else:  # shared batch dimensions -> einsum
+        result_shape = merge_shapes(x.shape.without(x_dims), y.shape.without(y_dims))
         REDUCE_LETTERS = list('ijklmn')
         KEEP_LETTERS = list('abcdefgh')
         x_letters = [(REDUCE_LETTERS if dim in x_dims else KEEP_LETTERS).pop(0) for dim in x.shape.names]
-        x_letter_map = {dim: letter for dim, letter in zip(x.shape.names, x_letters)}
+        letter_map = {dim: letter for dim, letter in zip(x.shape.names, x_letters)}
         REDUCE_LETTERS = list('ijklmn')
         y_letters = []
         for dim in y.shape.names:
@@ -1418,13 +1419,14 @@ def dot(x: Tensor,
                 y_letters.append(REDUCE_LETTERS.pop(0))
             else:
                 if dim in x.shape and dim not in x_dims:
-                    y_letters.append(x_letter_map[dim])
+                    y_letters.append(letter_map[dim])
                 else:
-                    y_letters.append(KEEP_LETTERS.pop(0))
-        keep_letters = list('abcdefgh')[:-len(KEEP_LETTERS)]
+                    next_letter = KEEP_LETTERS.pop(0)
+                    letter_map[dim] = next_letter
+                    y_letters.append(next_letter)
+        keep_letters = [letter_map[dim] for dim in result_shape.names]
         subscripts = f'{"".join(x_letters)},{"".join(y_letters)}->{"".join(keep_letters)}'
         result_native = backend.einsum(subscripts, x_native, y_native)
-        result_shape = merge_shapes(x.shape.without(x_dims), y.shape.without(y_dims))  # don't check group match  ToDo the order might be incorrect here
     return NativeTensor(result_native, result_shape)
 
 

From 96f22288a4f4504ef75414e5615bf1713dafd8d2 Mon Sep 17 00:00:00 2001
From: Philipp Holl <philipp.holl@tum.de>
Date: Tue, 15 Nov 2022 11:20:31 +0100
Subject: [PATCH 06/29] [math] Do not flatten batch dimensions by default.

---
 phi/math/_magic_ops.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/phi/math/_magic_ops.py b/phi/math/_magic_ops.py
index c9157c72e..6e76744dc 100644
--- a/phi/math/_magic_ops.py
+++ b/phi/math/_magic_ops.py
@@ -406,7 +406,7 @@ def unpack_dim(value, dim: str or Shape, unpacked_dims: Shape, **kwargs):
 
 
 
-def flatten(value, flat_dim: Shape = instance('flat'), **kwargs):
+def flatten(value, flat_dim: Shape = instance('flat'), flatten_batch=False, **kwargs):
     """
     Returns a `Tensor` with the same values as `value` but only a single dimension `flat_dim`.
     The order of the values in memory is not changed.
@@ -414,6 +414,8 @@ def flatten(value, flat_dim: Shape = instance('flat'), **kwargs):
     Args:
         value: `phi.math.magic.Shapable`, such as `Tensor`.
         flat_dim: Dimension name and type as `Shape` object. The size is ignored.
+        flatten_batch: Whether to flatten batch dimensions as well.
+            If `False`, batch dimensions are kept, only onn-batch dimensions are flattened.
         **kwargs: Additional keyword arguments required by specific implementations.
             Adding spatial dimensions to fields requires the `bounds: Box` argument specifying the physical extent of the new dimensions.
             Adding batch dimensions must always work without keyword arguments.
@@ -434,7 +436,7 @@ def flatten(value, flat_dim: Shape = instance('flat'), **kwargs):
         if result is not NotImplemented:
             return result
     # Fallback: pack_dims
-    return pack_dims(value, shape(value), flat_dim, **kwargs)
+    return pack_dims(value, shape(value) if flatten_batch else non_batch(value), flat_dim, **kwargs)
 
 
 # PhiTreeNode

From b946fdd765454ae5d484ead97a8f88bd52fa1481 Mon Sep 17 00:00:00 2001
From: Philipp Holl <philipp.holl@tum.de>
Date: Fri, 2 Dec 2022 11:47:53 -0600
Subject: [PATCH 07/29] [math] Fix lost item names in stack()

---
 phi/math/_shape.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/phi/math/_shape.py b/phi/math/_shape.py
index 3d8a93923..d91ca8d75 100644
--- a/phi/math/_shape.py
+++ b/phi/math/_shape.py
@@ -758,7 +758,7 @@ def _size_and_item_names_from_obj(obj, prev_size, prev_item_names, keep_item_nam
         if isinstance(obj, (tuple, list)):
             return len(obj), tuple(obj)
         elif isinstance(obj, Number):
-            return obj, prev_item_names if keep_item_names and obj == prev_size else None
+            return obj, prev_item_names if keep_item_names and (obj == prev_size or prev_size is None) else None
         elif isinstance(obj, math.Tensor) or obj is None:
             return obj, None
         else:

From b74b6191d2527419b571c1d7bc6356210784dadc Mon Sep 17 00:00:00 2001
From: Philipp Holl <philipp.holl@tum.de>
Date: Sat, 3 Dec 2022 14:53:40 -0600
Subject: [PATCH 08/29] [math] Fix tensor-Shape operations

---
 phi/math/_magic_ops.py | 8 ++++----
 phi/math/_tensors.py   | 6 +++---
 2 files changed, 7 insertions(+), 7 deletions(-)

diff --git a/phi/math/_magic_ops.py b/phi/math/_magic_ops.py
index 6e76744dc..d19727483 100644
--- a/phi/math/_magic_ops.py
+++ b/phi/math/_magic_ops.py
@@ -248,7 +248,7 @@ def expand(value, dims: Shape, **kwargs):
     # Fallback: stack
     if hasattr(value, '__stack__'):
         if dims.volume > 8:
-            warnings.warn(f"expand() default implementation is slow on large shapes {dims}. Please implement __expand__()", RuntimeWarning, stacklevel=2)
+            warnings.warn(f"expand() default implementation is slow on large shapes {dims}. Please implement __expand__() for {type(value).__name__} as defined in phi.math.magic", RuntimeWarning, stacklevel=2)
         for dim in reversed(dims):
             value = stack((value,) * dim.size, dim, **kwargs)
             assert value is not NotImplemented, "Value must implement either __expand__ or __stack__"
@@ -295,7 +295,7 @@ def rename_dims(value,
             return result
     # Fallback: unstack and stack
     if shape(value).only(dims).volume > 8:
-        warnings.warn(f"rename_dims() default implementation is slow on large dimensions ({shape(value).only(dims)}). Please implement __replace_dims__()", RuntimeWarning, stacklevel=2)
+        warnings.warn(f"rename_dims() default implementation is slow on large dimensions ({shape(value).only(dims)}). Please implement __replace_dims__() for {type(value).__name__} as defined in phi.math.magic", RuntimeWarning, stacklevel=2)
     for old_name, new_dim in zip(dims.names, names):
         value = stack(unstack(value, old_name), new_dim, **kwargs)
     return value
@@ -352,7 +352,7 @@ def pack_dims(value, dims: DimFilter, packed_dim: Shape, pos: int or None = None
             return result
     # Fallback: unstack and stack
     if shape(value).only(dims).volume > 8:
-        warnings.warn(f"pack_dims() default implementation is slow on large dimensions ({shape(value).only(dims)}). Please implement __pack_dims__()", RuntimeWarning, stacklevel=2)
+        warnings.warn(f"pack_dims() default implementation is slow on large dimensions ({shape(value).only(dims)}). Please implement __pack_dims__() for {type(value).__name__} as defined in phi.math.magic", RuntimeWarning, stacklevel=2)
     return stack(unstack(value, dims), packed_dim, **kwargs)
 
 
@@ -398,7 +398,7 @@ def unpack_dim(value, dim: str or Shape, unpacked_dims: Shape, **kwargs):
             return result
     # Fallback: unstack and stack
     if shape(value).only(dim).volume > 8:
-        warnings.warn(f"pack_dims() default implementation is slow on large dimensions ({shape(value).only(dim)}). Please implement __unpack_dim__()", RuntimeWarning, stacklevel=2)
+        warnings.warn(f"pack_dims() default implementation is slow on large dimensions ({shape(value).only(dim)}). Please implement __unpack_dim__() for {type(value).__name__} as defined in phi.math.magic", RuntimeWarning, stacklevel=2)
     unstacked = unstack(value, dim)
     for dim in reversed(unpacked_dims):
         unstacked = [stack(unstacked[i:i+dim.size], dim, **kwargs) for i in range(0, len(unstacked), dim.size)]
diff --git a/phi/math/_tensors.py b/phi/math/_tensors.py
index a77811280..c175aa3a7 100644
--- a/phi/math/_tensors.py
+++ b/phi/math/_tensors.py
@@ -1148,11 +1148,11 @@ def _op1(self, native_function):
 
     def _op2(self, other, operator, native_function, op_name: str = 'unknown', op_symbol: str = '?'):
         try:
-            other = self._tensor(other)
+            other_tensor = self._tensor(other)
         except NoBackendFound:
             return NotImplemented
-        if isinstance(other, NativeTensor):
-            return op2_native(self, other, native_function)
+        if isinstance(other_tensor, NativeTensor) or (isinstance(other_tensor, Tensor) and not isinstance(other, Tensor)):
+            return op2_native(self, other_tensor, native_function)
         else:
             return NotImplemented
 

From f0fa5e4f5fd91240a806b7d1ad9cf3e59551f849 Mon Sep 17 00:00:00 2001
From: Philipp Holl <philipp.holl@tum.de>
Date: Mon, 5 Dec 2022 21:02:24 +0100
Subject: [PATCH 09/29] [field] Fix resolution_from_staggered_tensor

---
 phi/field/_grid.py | 16 +++++++++++++---
 1 file changed, 13 insertions(+), 3 deletions(-)

diff --git a/phi/field/_grid.py b/phi/field/_grid.py
index 5b90af34e..657b71d1a 100644
--- a/phi/field/_grid.py
+++ b/phi/field/_grid.py
@@ -1,4 +1,4 @@
-from typing import TypeVar, Any
+from typing import TypeVar, Any, Tuple
 
 from phi import math, geom
 from phi.geom import Box, Geometry, GridCell
@@ -70,6 +70,16 @@ def __value_attrs__(self):
     def __variable_attrs__(self):
         return '_values',
 
+    def __expand__(self, dims: Shape, **kwargs) -> 'Grid':
+        return self.with_values(math.expand(self.values, dims, **kwargs))
+
+    def __replace_dims__(self, dims: Tuple[str, ...], new_dims: Shape, **kwargs) -> 'Grid':
+        for dim in dims:
+            assert dim not in self.resolution
+        values = math.rename_dims(self.values, dims, new_dims)
+        return self.with_values(values)
+
+
     def __eq__(self, other):
         if not type(self) == type(other):
             return False
@@ -487,10 +497,10 @@ def expand_staggered(values: Tensor, resolution: Shape, extrapolation: Extrapola
 
 def resolution_from_staggered_tensor(values: Tensor, extrapolation: Extrapolation):
     any_dim = values.shape.spatial.names[0]
-    x = values.vector[any_dim]
+    x_shape = values.shape.after_gather({'vector': any_dim})
     ext_lower, ext_upper = extrapolation.valid_outer_faces(any_dim)
     delta = int(ext_lower) + int(ext_upper) - 1
-    resolution = x.shape.spatial._replace_single_size(any_dim, x.shape.get_size(any_dim) - delta)
+    resolution = x_shape.spatial._replace_single_size(any_dim, x_shape.get_size(any_dim) - delta)
     return resolution
 
 

From f7f2326e7f5d7f70bc1d5fac8eb4e54fb8a6801c Mon Sep 17 00:00:00 2001
From: Philipp Holl <philipp.holl@tum.de>
Date: Tue, 6 Dec 2022 15:46:19 +0100
Subject: [PATCH 10/29] [math] Only show conversion warning when rank > 1

---
 phi/math/_tensors.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/phi/math/_tensors.py b/phi/math/_tensors.py
index c175aa3a7..aece1641b 100644
--- a/phi/math/_tensors.py
+++ b/phi/math/_tensors.py
@@ -85,7 +85,8 @@ def numpy(self, order: str or tuple or list or Shape = None) -> np.ndarray:
         return choose_backend(native).numpy(native)
 
     def __array__(self, dtype=None):  # NumPy conversion
-        warnings.warn("Automatic conversion of Φ-Flow tensors to NumPy can cause problems because the dimension order is not guaranteed.", SyntaxWarning, stacklevel=3)
+        if self.rank > 1:
+            warnings.warn("Automatic conversion of Φ-Flow tensors to NumPy can cause problems because the dimension order is not guaranteed.", SyntaxWarning, stacklevel=3)
         return self.numpy(self._shape)
 
     def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):  # NumPy interface

From 0f36aa803d66e69e36b02dfdde33eea89838d051 Mon Sep 17 00:00:00 2001
From: Philipp Holl <philipp.holl@tum.de>
Date: Tue, 6 Dec 2022 16:04:45 +0100
Subject: [PATCH 11/29] [math] Make extrapolations Shapable

* Add unit test
---
 phi/field/_grid.py                      |  7 +++++--
 phi/math/_magic_ops.py                  |  2 +-
 phi/math/extrapolation.py               | 27 ++++++++++++++++++++++++-
 tests/commit/math/test_extrapolation.py | 13 +++++++++++-
 4 files changed, 44 insertions(+), 5 deletions(-)

diff --git a/phi/field/_grid.py b/phi/field/_grid.py
index 657b71d1a..8e7286c4a 100644
--- a/phi/field/_grid.py
+++ b/phi/field/_grid.py
@@ -75,9 +75,12 @@ def __expand__(self, dims: Shape, **kwargs) -> 'Grid':
 
     def __replace_dims__(self, dims: Tuple[str, ...], new_dims: Shape, **kwargs) -> 'Grid':
         for dim in dims:
-            assert dim not in self.resolution
+            if dim in self._resolution:
+                return NotImplemented
         values = math.rename_dims(self.values, dims, new_dims)
-        return self.with_values(values)
+        extrapolation = math.rename_dims(self.extrapolation, dims, new_dims, **kwargs)
+        bounds = math.rename_dims(self.bounds, dims, new_dims, **kwargs)
+        return type(self)(values, extrapolation=extrapolation, bounds=bounds, resolution=self._resolution)
 
 
     def __eq__(self, other):
diff --git a/phi/math/_magic_ops.py b/phi/math/_magic_ops.py
index d19727483..5b3ef4e38 100644
--- a/phi/math/_magic_ops.py
+++ b/phi/math/_magic_ops.py
@@ -286,7 +286,7 @@ def rename_dims(value,
     """
     if isinstance(value, Shape):
         return value._replace_names_and_types(dims, names)
-    assert isinstance(value, Shapable) and isinstance(value, Shaped), "value must be a Shape or Shapable."
+    assert isinstance(value, Shapable) and isinstance(value, Shaped), f"value must be a Shape or Shapable but got {type(value).__name__}"
     dims = shape(value).only(dims)
     names = dims._replace_names_and_types(dims, names)
     if hasattr(value, '__replace_dims__'):
diff --git a/phi/math/extrapolation.py b/phi/math/extrapolation.py
index 652afeb18..4e819b32f 100644
--- a/phi/math/extrapolation.py
+++ b/phi/math/extrapolation.py
@@ -10,7 +10,7 @@
 
 from phi.math.backend._backend import get_spatial_derivative_order
 from .backend import choose_backend
-from ._shape import Shape, channel, spatial
+from ._shape import Shape, channel, spatial, EMPTY_SHAPE, merge_shapes
 from ._magic_ops import concat, stack
 from ._tensors import Tensor, NativeTensor, CollapsedTensor, TensorStack, wrap
 from . import _ops as math  # TODO this executes _ops.py, can we avoid this?
@@ -196,6 +196,10 @@ def __init__(self, value: Tensor or float):
         self.value = wrap(value)
         """ Extrapolation value """
 
+    @property
+    def shape(self):
+        return self.value.shape
+
     def __repr__(self):
         return repr(self.value)
 
@@ -364,6 +368,10 @@ def __neg__(self):
 
 class _CopyExtrapolation(Extrapolation):
 
+    @property
+    def shape(self):
+        return EMPTY_SHAPE
+
     def is_copy_pad(self, dim: str, upper_edge: bool):
         return True
 
@@ -625,6 +633,11 @@ def transform_coordinates(self, coordinates: Tensor, shape: Shape, **kwargs) ->
 
 
 class _NoExtrapolation(Extrapolation):  # singleton
+
+    @property
+    def shape(self):
+        return EMPTY_SHAPE
+
     def to_dict(self) -> dict:
         return {'type': 'none'}
 
@@ -691,6 +704,10 @@ def __init__(self, derived_from: Extrapolation):
         super().__init__(-1)
         self.derived_from = derived_from
 
+    @property
+    def shape(self):
+        return EMPTY_SHAPE
+
     def to_dict(self) -> dict:
         return {'type': 'undefined', 'derived_from': self.derived_from.to_dict()}
 
@@ -806,6 +823,10 @@ def __init__(self, extrapolations: Dict[str, Tuple[Extrapolation, Extrapolation]
         super().__init__(pad_rank=None)
         self.ext = extrapolations
 
+    @property
+    def shape(self):
+        return merge_shapes(*sum(self.ext.values(), ()))
+
     def to_dict(self) -> dict:
         return {
             'type': 'mixed',
@@ -945,6 +966,10 @@ def __init__(self, normal: Extrapolation, tangential: Extrapolation):
         self.normal = normal
         self.tangential = tangential
 
+    @property
+    def shape(self):
+        return merge_shapes(self.normal, self.tangential)
+
     def to_dict(self) -> dict:
         return {
             'type': 'normal-tangential',
diff --git a/tests/commit/math/test_extrapolation.py b/tests/commit/math/test_extrapolation.py
index 5852f3e29..a8f4e27d9 100644
--- a/tests/commit/math/test_extrapolation.py
+++ b/tests/commit/math/test_extrapolation.py
@@ -1,7 +1,7 @@
 from unittest import TestCase
 
 import phi
-from phi.math import NUMPY, spatial, batch, extrapolation
+from phi.math import NUMPY, spatial, batch, extrapolation, shape
 from phi.math.extrapolation import *
 from phi import math
 
@@ -207,3 +207,14 @@ def test_slice_normal_tangential(self):
         ext = combine_sides(x=(INFLOW_LEFT, BOUNDARY), y=0)
         self.assertEqual(combine_sides(x=(1, BOUNDARY), y=0), ext[{'vector': 'x'}])
         self.assertEqual(combine_sides(x=(0, BOUNDARY), y=0), ext[{'vector': 'y'}])
+
+    def test_shapes(self):
+        self.assertEqual(EMPTY_SHAPE, ONE.shape)
+        self.assertEqual(EMPTY_SHAPE, PERIODIC.shape)
+        self.assertEqual(EMPTY_SHAPE, BOUNDARY.shape)
+        self.assertEqual(EMPTY_SHAPE, SYMMETRIC.shape)
+        self.assertEqual(EMPTY_SHAPE, REFLECT.shape)
+        v = math.vec(x=1, y=0)
+        self.assertEqual(v.shape, shape(ZERO + v))
+        self.assertEqual(v.shape, shape(combine_sides(x=v, y=0)))
+        self.assertEqual(v.shape, shape(combine_by_direction(normal=v, tangential=0)))

From 5788543238bf5f0f895a40991de5768c5778a42e Mon Sep 17 00:00:00 2001
From: Philipp Holl <philipp.holl@tum.de>
Date: Tue, 6 Dec 2022 16:09:00 +0100
Subject: [PATCH 12/29] [math] Fix iter for multi-dim tensors

---
 phi/math/_tensors.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/phi/math/_tensors.py b/phi/math/_tensors.py
index aece1641b..7dc113e58 100644
--- a/phi/math/_tensors.py
+++ b/phi/math/_tensors.py
@@ -664,8 +664,9 @@ def __iter__(self):
         elif self.rank == 0:
             return iter([self.native()])
         else:
-            from ._magic_ops import flatten
-            return iter(flatten(self))
+            from ._ops import reshaped_native
+            native = reshaped_native(self, [self.shape])
+            return iter(native)
 
     def _tensor(self, other):
         if isinstance(other, Tensor):

From e4ebe29e3901d66a4030c0e16d78410d5e4e5ea9 Mon Sep 17 00:00:00 2001
From: holl- <philipp@mholl.de>
Date: Mon, 12 Dec 2022 15:39:40 +0100
Subject: [PATCH 13/29] [math] Add flatten_batch argument to magic API

---
 phi/math/_magic_ops.py | 7 ++++---
 phi/math/_tensors.py   | 4 +++-
 phi/math/magic.py      | 5 +++--
 3 files changed, 10 insertions(+), 6 deletions(-)

diff --git a/phi/math/_magic_ops.py b/phi/math/_magic_ops.py
index 5b3ef4e38..e11ad02bd 100644
--- a/phi/math/_magic_ops.py
+++ b/phi/math/_magic_ops.py
@@ -387,7 +387,9 @@ def unpack_dim(value, dim: str or Shape, unpacked_dims: Shape, **kwargs):
     assert isinstance(value, Shapable) and isinstance(value, Sliceable) and isinstance(value, Shaped), f"value must be Shapable but got {type(value)}"
     if isinstance(dim, Shape):
         dim = dim.name
-    assert isinstance(dim, str), f"dim must be a str but got {type(dim)}"
+    assert isinstance(dim, str), f"dim must be a str or Shape but got {type(dim)}"
+    if dim not in shape(value):
+        return value  # Nothing to do, maybe expand?
     if unpacked_dims.rank == 0:
         return value[{dim: 0}]  # remove dim
     elif unpacked_dims.rank == 1:
@@ -405,7 +407,6 @@ def unpack_dim(value, dim: str or Shape, unpacked_dims: Shape, **kwargs):
     return unstacked[0]
 
 
-
 def flatten(value, flat_dim: Shape = instance('flat'), flatten_batch=False, **kwargs):
     """
     Returns a `Tensor` with the same values as `value` but only a single dimension `flat_dim`.
@@ -432,7 +433,7 @@ def flatten(value, flat_dim: Shape = instance('flat'), flatten_batch=False, **kw
     assert isinstance(flat_dim, Shape) and flat_dim.rank == 1, flat_dim
     assert isinstance(value, Shapable) and isinstance(value, Shaped), f"value must be Shapable but got {type(value)}"
     if hasattr(value, '__flatten__'):
-        result = value.__flatten__(flat_dim, **kwargs)
+        result = value.__flatten__(flat_dim, flatten_batch, **kwargs)
         if result is not NotImplemented:
             return result
     # Fallback: pack_dims
diff --git a/phi/math/_tensors.py b/phi/math/_tensors.py
index 7dc113e58..8d9366d09 100644
--- a/phi/math/_tensors.py
+++ b/phi/math/_tensors.py
@@ -923,7 +923,9 @@ def __stack__(self, values: tuple, dim: Shape, **kwargs) -> 'Shapable':
     def __concat__(self, values: tuple, dim: str, **kwargs) -> 'Shapable':
         return NotImplemented
 
-    def __flatten__(self, flat_dim: Shape):
+    def __flatten__(self, flat_dim: Shape, flatten_batch: bool):
+        if not flatten_batch and self._shape.batch:
+            raise NotImplementedError
         return layout(self._as_list(), flat_dim)
 
     def __expand__(self, dims: Shape, **kwargs) -> 'Tensor':
diff --git a/phi/math/magic.py b/phi/math/magic.py
index e2ec74637..a9fb23736 100644
--- a/phi/math/magic.py
+++ b/phi/math/magic.py
@@ -272,13 +272,15 @@ def __unpack_dim__(self, dim: str, unpacked_dims: Shape, **kwargs) -> 'Shapable'
         """
         raise NotImplementedError
 
-    def __flatten__(self, flat_dim: Shape, **kwargs):
+    def __flatten__(self, flat_dim: Shape, flatten_batch: bool, **kwargs) -> 'Shapable':
         """
         Lays out all elements along a single dimension.
         This is equivalent to packing all dimensions.
 
         Args:
             flat_dim: Single dimension as `Shape`.
+            flatten_batch: Whether to flatten batch dimensions as well.
+            If `False`, batch dimensions are kept, only onn-batch dimensions are flattened.
             **kwargs: Additional keyword arguments required by specific implementations.
                 Adding spatial dimensions to fields requires the `bounds: Box` argument specifying the physical extent of the new dimensions.
                 Adding batch dimensions must always work without keyword arguments.
@@ -289,7 +291,6 @@ def __flatten__(self, flat_dim: Shape, **kwargs):
         raise NotImplementedError
 
 
-
 class _PhiTreeNodeType(type):
 
     def __instancecheck__(self, instance):

From 86cc88d62530ab8591c4f68f19009201ec907121 Mon Sep 17 00:00:00 2001
From: holl- <philipp@mholl.de>
Date: Mon, 12 Dec 2022 15:39:56 +0100
Subject: [PATCH 14/29] [geom] Fast reshaping operations for Sphere and Point

---
 phi/geom/_geom.py                 | 25 +++++++++
 phi/geom/_sphere.py               | 29 +++++++++--
 tests/commit/geom/test__sphere.py | 87 +++++++++----------------------
 tests/commit/geom/test_geom.py    | 12 +++++
 4 files changed, 87 insertions(+), 66 deletions(-)

diff --git a/phi/geom/_geom.py b/phi/geom/_geom.py
index bbb373bc7..e65f53692 100644
--- a/phi/geom/_geom.py
+++ b/phi/geom/_geom.py
@@ -1,4 +1,5 @@
 from numbers import Number
+from typing import Tuple
 
 import numpy as np
 
@@ -515,6 +516,30 @@ def __stack__(self, values: tuple, dim: Shape, **kwargs) -> 'Geometry':
         else:
             return Geometry.__stack__(self, values, dim, **kwargs)
 
+    def __concat__(self, values: tuple, dim: str, **kwargs) -> 'Point':
+        if all(isinstance(v, Point) for v in values):
+            return Point(math.concat([v.center for v in values], dim, **kwargs))
+        else:
+            return NotImplemented
+
+    def __replace_dims__(self, dims: Tuple[str, ...], new_dims: Shape, **kwargs) -> 'Point':
+        return Point(math.rename_dims(self._location, dims, new_dims, **kwargs))
+
+    def __expand__(self, dims: Shape, **kwargs) -> 'Point':
+        return Point(math.expand(self._location, dims, **kwargs))
+
+    def __pack_dims__(self, dims: Tuple[str, ...], packed_dim: Shape, pos: int or None, **kwargs) -> 'Point':
+        return Point(math.pack_dims(self._location, dims, packed_dim, pos, **kwargs))
+
+    def __unpack_dim__(self, dim: str, unpacked_dims: Shape, **kwargs) -> 'Point':
+        return Point(math.unpack_dim(self._location, dim, unpacked_dims, **kwargs))
+
+    def __flatten__(self, flat_dim: Shape, flatten_batch: bool, **kwargs) -> 'Point':
+        dims = self.shape.without('vector')
+        if not flatten_batch:
+            dims = dims.non_batch
+        return Point(math.pack_dims(self._location, dims, flat_dim, **kwargs))
+
 
 def assert_same_rank(rank1, rank2, error_message):
     """ Tests that two objects have the same spatial rank. Objects can be of types: `int`, `None` (no check), `Geometry`, `Shape`, `Tensor` """
diff --git a/phi/geom/_sphere.py b/phi/geom/_sphere.py
index b7d8d2a73..7ee000959 100644
--- a/phi/geom/_sphere.py
+++ b/phi/geom/_sphere.py
@@ -1,11 +1,8 @@
-import warnings
-from typing import Dict
+from typing import Tuple
 
 from phi import math
-
 from ._geom import Geometry, _keep_vector
 from ..math import wrap, Tensor, Shape
-from ..math.backend import PHI_LOGGER
 from ..math.magic import slicing_dict
 
 
@@ -120,6 +117,30 @@ def __stack__(self, values: tuple, dim: Shape, **kwargs) -> 'Geometry':
         else:
             return Geometry.__stack__(self, values, dim, **kwargs)
 
+    def __concat__(self, values: tuple, dim: str, **kwargs) -> 'Sphere':
+        if all(isinstance(v, Sphere) for v in values):
+            return Sphere(math.concat([v.center for v in values], dim, **kwargs), radius=math.concat([v.radius for v in values], dim, **kwargs))
+        else:
+            return NotImplemented
+
+    def __replace_dims__(self, dims: Tuple[str, ...], new_dims: Shape, **kwargs) -> 'Sphere':
+        return Sphere(math.rename_dims(self._center, dims, new_dims, **kwargs), radius=math.rename_dims(self._radius, dims, new_dims, **kwargs))
+
+    def __expand__(self, dims: Shape, **kwargs) -> 'Sphere':
+        return Sphere(math.expand(self._center, dims, **kwargs), radius=self._radius)
+
+    def __pack_dims__(self, dims: Tuple[str, ...], packed_dim: Shape, pos: int or None, **kwargs) -> 'Sphere':
+        return Sphere(math.pack_dims(self._center, dims, packed_dim, pos, **kwargs), radius=math.pack_dims(self._radius, dims, packed_dim, pos, **kwargs))
+
+    def __unpack_dim__(self, dim: str, unpacked_dims: Shape, **kwargs) -> 'Sphere':
+        return Sphere(math.unpack_dim(self._center, dim, unpacked_dims, **kwargs), radius=math.unpack_dim(self._radius, dim, unpacked_dims, **kwargs))
+
+    def __flatten__(self, flat_dim: Shape, flatten_batch: bool, **kwargs) -> 'Shapable':
+        dims = self.shape.without('vector')
+        if not flatten_batch:
+            dims = dims.non_batch
+        return Sphere(math.pack_dims(self._center, dims, flat_dim, **kwargs), radius=math.pack_dims(self._radius, dims, flat_dim, **kwargs))
+
     def push(self, positions: Tensor, outward: bool = True, shift_amount: float = 0) -> Tensor:
         raise NotImplementedError()
 
diff --git a/tests/commit/geom/test__sphere.py b/tests/commit/geom/test__sphere.py
index b6d6cfd4f..02339a340 100644
--- a/tests/commit/geom/test__sphere.py
+++ b/tests/commit/geom/test__sphere.py
@@ -1,9 +1,8 @@
 from unittest import TestCase
 
-
 from phi import math
-from phi.geom import Box, union, Cuboid, embed, Sphere
-from phi.math import batch, channel
+from phi.geom import union, Sphere
+from phi.math import stack, vec, instance, expand, rename_dims, unpack_dim, pack_dims, spatial, flatten, batch, channel
 from phi.math.magic import Shaped, Sliceable, Shapable
 
 
@@ -43,62 +42,26 @@ def test_project(self):
         self.assertEqual(sphere, sphere.vector['x,y'])
         self.assertEqual(Sphere(x=4, radius=1), sphere['x'])
 
-    # def test_box_constructor(self):
-    #     try:
-    #         Box(0, (1, 1))
-    #         raise RuntimeError
-    #     except AssertionError:
-    #         pass
-    #     math.assert_close(Box(x=1, y=1).size, 1)
-    #
-    # def test_box_batched(self):
-    #     lower = math.tensor([(0, 0), (1, 1)], batch('boxes'), channel(vector='x,y'))
-    #     upper = math.wrap((1, 1), channel(vector='x,y'))
-    #     box = Box(lower, upper)
-    #     self.assertEqual(batch(boxes=2) & channel(vector='x,y'), box.shape)
-    #
-    # def test_slice(self):
-    #     b1, b2 = Box(x=4, y=3), Box(x=2, y=1)
-    #     u = union(b1, b2)
-    #     self.assertEqual(b1, u.union[0])
-    #     self.assertEqual(b2, u.union[1])
-    #
-    # def test_without(self):
-    #     box = Box(x=4, y=3)
-    #     self.assertEqual(Box(x=4), box.without(('y',)))
-    #     self.assertEqual(Box(), box.without(('x', 'y')))
-    #     self.assertEqual(box, box.without(()))
-    #
-    # def test_embed(self):
-    #     self.assertEqual(Box(x=4, y=3, z=None), embed(Box(x=4, y=3), 'z'))
-    #     self.assertEqual(Box(x=4, y=3, z=None), embed(Box(x=4, y=3), 'x,z'))
-    #
-    # def test_box_product(self):
-    #     a = Box(x=4)
-    #     b = Box(y=3).shifted(math.wrap(1))
-    #     ab = a * b
-    #     self.assertEqual(2, ab.spatial_rank)
-    #     math.assert_close(ab.size, (4, 3))
-    #     math.assert_close(ab.lower, (0, 1))
-    #
-    # def test_union_same(self):
-    #     u = union(Box(x=1, y=1), Box(x=(2, 3), y=1))
-    #     self.assertIsInstance(u, Box)
-    #     math.assert_close(u.approximate_signed_distance((0, 0)), u.approximate_signed_distance((3, 1)), 0)
-    #     math.assert_close(u.approximate_signed_distance((1.5, 0)), 0.5)
-    #
-    # def test_stack_volume(self):
-    #     u = math.stack([Box(x=1, y=1), Box(x=2, y=2)], batch('batch'))
-    #     math.assert_close(u.volume, [1, 4])
-    #
-    # def test_shape_type(self):
-    #     box = Box(x=1, y=2)
-    #     self.assertEqual(box.rotated(0.1).shape_type, 'rotB')
-    #
-    # def test_box_eq(self):
-    #     self.assertNotEqual(Box(x=1, y=1), Box(x=1))
-    #     self.assertEqual(Box(x=1, y=1), Box(x=1, y=1))
-    #
-    # def test_cuboid_constructor_kwargs(self):
-    #     c = Cuboid(x=2., y=1.)
-    #     math.assert_close(c.lower, -c.upper, (-1, -.5))
+    def test_reshaping(self):
+        s = stack([Sphere(vec(x=0, y=0), radius=1)] * 50, instance('points'))
+        s = expand(s, batch(b=100))
+        s = rename_dims(s, 'b', 'bat')
+        s = unpack_dim(s, 'points', spatial(x=10, y=5))
+        assert batch(bat=100) & spatial(x=10, y=5) & channel(vector='x,y') == s.shape
+        s = pack_dims(s, 'x,y', instance('particles'))
+        assert batch(bat=100) & instance(particles=50) & channel(vector='x,y') == s.shape
+        s = flatten(s)
+        assert batch(bat=100) & instance(flat=50) & channel(vector='x,y') == s.shape
+
+    def test_reshaping_const_radius(self):
+        s = Sphere(stack([vec(x=0, y=0)] * 50, instance('points')), radius=1)
+        s = expand(s, batch(b=100))
+        s = rename_dims(s, 'b', 'bat')
+        s = unpack_dim(s, 'points', spatial(x=10, y=5))
+        assert not s.radius.shape
+        assert batch(bat=100) & spatial(x=10, y=5) & channel(vector='x,y') == s.shape
+        s = pack_dims(s, 'x,y', instance('particles'))
+        assert not s.radius.shape
+        assert batch(bat=100) & instance(particles=50) & channel(vector='x,y') == s.shape
+        s = flatten(s)
+        assert batch(bat=100) & instance(flat=50) & channel(vector='x,y') == s.shape
diff --git a/tests/commit/geom/test_geom.py b/tests/commit/geom/test_geom.py
index c3f4695c1..97558dcd1 100644
--- a/tests/commit/geom/test_geom.py
+++ b/tests/commit/geom/test_geom.py
@@ -1,6 +1,7 @@
 from unittest import TestCase
 
 from phi import math, geom
+from phi.math import stack, vec, instance, expand, rename_dims, unpack_dim, pack_dims, spatial, flatten, batch, channel
 from phi.geom import Box, Sphere
 
 
@@ -24,3 +25,14 @@ def test_infinite_cylinder(self):
         corner_distance = math.sqrt(2) / 2 - .5
         distance = math.wrap([corner_distance, 0, corner_distance, corner_distance, 0], math.instance('points'))
         math.assert_close(cylinder.approximate_signed_distance(loc), distance)
+
+    def test_point_reshaping(self):
+        s = stack([geom.Point(vec(x=0, y=0))] * 50, instance('points'))
+        s = expand(s, batch(b=100))
+        s = rename_dims(s, 'b', 'bat')
+        s = unpack_dim(s, 'points', spatial(x=10, y=5))
+        assert batch(bat=100) & spatial(x=10, y=5) & channel(vector='x,y') == s.shape
+        s = pack_dims(s, 'x,y', instance('particles'))
+        assert batch(bat=100) & instance(particles=50) & channel(vector='x,y') == s.shape
+        s = flatten(s)
+        assert batch(bat=100) & instance(flat=50) & channel(vector='x,y') == s.shape
\ No newline at end of file

From 09140eb093f22d2d7b9f56a4e22fab90084fc4f7 Mon Sep 17 00:00:00 2001
From: holl- <philipp@mholl.de>
Date: Mon, 12 Dec 2022 15:40:30 +0100
Subject: [PATCH 15/29] [docs] Fix diffuse.explicit() docstring

---
 phi/physics/diffuse.py | 3 ---
 1 file changed, 3 deletions(-)

diff --git a/phi/physics/diffuse.py b/phi/physics/diffuse.py
index 00fbb49b9..a7ddd04d7 100644
--- a/phi/physics/diffuse.py
+++ b/phi/physics/diffuse.py
@@ -15,9 +15,6 @@ def explicit(field: FieldType,
     """
     Simulate a finite-time diffusion process of the form dF/dt = α · ΔF on a given `Field` FieldType with diffusion coefficient α.
 
-    If `field` is periodic (set via `extrapolation='periodic'`), diffusion may be simulated in Fourier space.
-    Otherwise, finite differencing is used to approximate the
-
     Args:
         field: CenteredGrid, StaggeredGrid or ConstantField
         diffusivity: Diffusion per time. `diffusion_amount = diffusivity * dt`

From 690c4f475e68d85702e6349f301b5e6b89a71580 Mon Sep 17 00:00:00 2001
From: holl- <philipp@mholl.de>
Date: Mon, 12 Dec 2022 15:43:29 +0100
Subject: [PATCH 16/29] [field] Fast reshaping for PointClouds

---
 phi/field/_point_cloud.py               | 11 ++++++++++-
 tests/commit/field/test__point_cloud.py | 14 ++++++++++++++
 2 files changed, 24 insertions(+), 1 deletion(-)
 create mode 100644 tests/commit/field/test__point_cloud.py

diff --git a/phi/field/_point_cloud.py b/phi/field/_point_cloud.py
index d16346a23..e20655dad 100644
--- a/phi/field/_point_cloud.py
+++ b/phi/field/_point_cloud.py
@@ -1,5 +1,5 @@
 import warnings
-from typing import Any
+from typing import Any, Tuple
 
 from phi import math
 from phi.geom import Geometry, GridCell, Box, Point
@@ -82,6 +82,15 @@ def __value_attrs__(self):
     def __variable_attrs__(self):
         return '_values', '_elements'
 
+    def __expand__(self, dims: Shape, **kwargs) -> 'PointCloud':
+        return self.with_values(math.expand(self.values, dims, **kwargs))
+
+    def __replace_dims__(self, dims: Tuple[str, ...], new_dims: Shape, **kwargs) -> 'PointCloud':
+        elements = math.rename_dims(self.elements, dims, new_dims)
+        values = math.rename_dims(self.values, dims, new_dims)
+        extrapolation = math.rename_dims(self.extrapolation, dims, new_dims, **kwargs)
+        return PointCloud(elements, values, extrapolation, self._add_overlapping, self._bounds, self._color)
+
     def __eq__(self, other):
         if not type(self) == type(other):
             return False
diff --git a/tests/commit/field/test__point_cloud.py b/tests/commit/field/test__point_cloud.py
new file mode 100644
index 000000000..f9fb889dc
--- /dev/null
+++ b/tests/commit/field/test__point_cloud.py
@@ -0,0 +1,14 @@
+from unittest import TestCase
+
+from phi.field import PointCloud
+from phi.geom import Sphere
+from phi.math import batch, stack, instance, expand, rename_dims, shape, vec
+
+
+class GridTest(TestCase):
+
+    def test_reshape(self):
+        c = PointCloud(Sphere(stack([vec(x=0, y=1)] * 50, instance('points')), radius=.1))
+        c = expand(c, batch(b=2))
+        c = rename_dims(c, 'points', 'particles')
+        assert batch(b=2) & instance(particles=50) == shape(c)

From 843fa7cd066c445ce5f4dfcd2b124bd0baf67e22 Mon Sep 17 00:00:00 2001
From: holl- <philipp@mholl.de>
Date: Mon, 12 Dec 2022 15:49:29 +0100
Subject: [PATCH 17/29] [math] Fix flatten usages

Most usages expect flatten to reduce batches as well. This is now handled explicitly but we may have to rollback the default to False.
---
 phi/field/_scene.py  | 14 +++++++-------
 phi/math/_ops.py     |  2 +-
 phi/math/_tensors.py |  2 +-
 3 files changed, 9 insertions(+), 9 deletions(-)

diff --git a/phi/field/_scene.py b/phi/field/_scene.py
index cc92999bb..242a80ffe 100644
--- a/phi/field/_scene.py
+++ b/phi/field/_scene.py
@@ -199,7 +199,7 @@ def at(directory: str or tuple or list or math.Tensor or 'Scene', id: int or mat
             id = math.wrap(id)
             paths = math.map(lambda d, i: join(d, f"sim_{i:06d}"), directory, id)
         # test all exist
-        for path in math.flatten(paths):
+        for path in math.flatten(paths, flatten_batch=True):
             if not isdir(path):
                 raise IOError(f"There is no scene at '{path}'")
         return Scene(paths)
@@ -266,7 +266,7 @@ def exist_properties(self):
         if self._properties is not None:
             return True  # must have been written or read
         else:
-            json_file = join(next(iter(math.flatten(self._paths))), "description.json")
+            json_file = join(next(iter(math.flatten(self._paths, flatten_batch=True))), "description.json")
             return isfile(json_file)
 
     def exists_config(self):
@@ -454,7 +454,7 @@ def copy_calling_script(self, full_trace=False, include_context_information=True
                 text = "\n\n".join(blocks)
                 self.copy_src_text('ipython.py', text)
         if include_context_information:
-            for path in math.flatten(self._paths):
+            for path in math.flatten(self._paths, flatten_batch=True):
                 with open(join(path, 'src', 'context.json'), 'w') as context_file:
                     json.dump({
                         'phi_version': phi_version,
@@ -462,23 +462,23 @@ def copy_calling_script(self, full_trace=False, include_context_information=True
                     }, context_file)
 
     def copy_src(self, script_path, only_external=True):
-        for path in math.flatten(self._paths):
+        for path in math.flatten(self._paths, flatten_batch=True):
             if not only_external or not _is_phi_file(script_path):
                 shutil.copy(script_path, join(path, 'src', basename(script_path)))
 
     def copy_src_text(self, filename, text):
-        for path in math.flatten(self._paths):
+        for path in math.flatten(self._paths, flatten_batch=True):
             target = join(path, 'src', filename)
             with open(target, "w") as file:
                 file.writelines(text)
 
     def mkdir(self):
-        for path in math.flatten(self._paths):
+        for path in math.flatten(self._paths, flatten_batch=True):
             isdir(path) or os.mkdir(path)
 
     def remove(self):
         """ Deletes the scene directory and all contained files. """
-        for p in math.flatten(self._paths):
+        for p in math.flatten(self._paths, flatten_batch=True):
             p = abspath(p)
             if isdir(p):
                 shutil.rmtree(p)
diff --git a/phi/math/_ops.py b/phi/math/_ops.py
index 07bde820b..0c76d7250 100644
--- a/phi/math/_ops.py
+++ b/phi/math/_ops.py
@@ -380,7 +380,7 @@ def map_(function, *values, **kwargs) -> Tensor or None:
     values = [wrap(v) for v in values]
     shape = merge_shapes(*[v.shape for v in values])
     values_reshaped = [expand(v, shape) for v in values]
-    flat = [flatten(v) for v in values_reshaped]
+    flat = [flatten(v, flatten_batch=True) for v in values_reshaped]
     result = []
     for items in zip(*flat):
         result.append(function(*items, **kwargs))
diff --git a/phi/math/_tensors.py b/phi/math/_tensors.py
index 8d9366d09..d247c1a74 100644
--- a/phi/math/_tensors.py
+++ b/phi/math/_tensors.py
@@ -1018,7 +1018,7 @@ def _tensor_reduce(self,
                 result = native_function(choose_backend(self._as_list()), self._obj, 0)
             return wrap(result)
         if not self._shape.without(dims):
-            return self.__flatten__(batch('_flat'))._tensor_reduce(('_flat',), dtype, native_function, collapsed_function, unaffected_function)
+            return self.__flatten__(batch('_flat'), flatten_batch=True)._tensor_reduce(('_flat',), dtype, native_function, collapsed_function, unaffected_function)
         else:
             raise NotImplementedError(f"Partial Layout reduction not yet supported. Shape={self._shape}, reduce={dims}")
         # # --- inner reduce ---

From a353832e1da35d092c593228e4b91e4256476162 Mon Sep 17 00:00:00 2001
From: holl- <philipp@mholl.de>
Date: Mon, 12 Dec 2022 16:23:51 +0100
Subject: [PATCH 18/29] [math] Fix bool bug in Shape

---
 phi/math/_shape.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/phi/math/_shape.py b/phi/math/_shape.py
index d91ca8d75..4e83e3f42 100644
--- a/phi/math/_shape.py
+++ b/phi/math/_shape.py
@@ -758,7 +758,7 @@ def _size_and_item_names_from_obj(obj, prev_size, prev_item_names, keep_item_nam
         if isinstance(obj, (tuple, list)):
             return len(obj), tuple(obj)
         elif isinstance(obj, Number):
-            return obj, prev_item_names if keep_item_names and (obj == prev_size or prev_size is None) else None
+            return obj, prev_item_names if keep_item_names and (prev_size is None or _size_equal(obj, prev_size)) else None
         elif isinstance(obj, math.Tensor) or obj is None:
             return obj, None
         else:

From e5873e442508b82bc7016da956ed6ec5f0cd3409 Mon Sep 17 00:00:00 2001
From: Philipp Holl <philipp.holl@tum.de>
Date: Mon, 12 Dec 2022 21:46:19 +0100
Subject: [PATCH 19/29] [tests] Add grid reshaping unit tests

---
 tests/commit/field/test__grid.py | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/tests/commit/field/test__grid.py b/tests/commit/field/test__grid.py
index 074e562ab..6865a8b79 100644
--- a/tests/commit/field/test__grid.py
+++ b/tests/commit/field/test__grid.py
@@ -138,3 +138,13 @@ def test_is_phi_tree_node(self):
         self.assertTrue(issubclass(CenteredGrid, PhiTreeNode))
         grid = CenteredGrid(0, x=4)
         self.assertTrue(isinstance(grid, PhiTreeNode))
+
+    def test_reshape_centered_grid(self):
+        grid = math.expand(CenteredGrid(1, x=10, y=10), batch(b=100))
+        grid = math.rename_dims(grid, 'b', 'bat')
+        self.assertEqual(batch(bat=100) & spatial(x=10, y=10), grid.shape)
+
+    def test_reshape_staggered_grid(self):
+        grid = math.expand(StaggeredGrid(1, x=10, y=10), batch(b=100))
+        grid = math.rename_dims(grid, 'b', 'bat')
+        self.assertEqual(batch(bat=100) & spatial(x=10, y=10) & channel(vector='x,y'), grid.shape)

From 521c604b569b474ce45010330b637bf9a7012b6e Mon Sep 17 00:00:00 2001
From: Philipp Holl <philipp.holl@tum.de>
Date: Mon, 12 Dec 2022 21:49:38 +0100
Subject: [PATCH 20/29] [ci] Don't test on Python 3.6 due to GitHub actions
 limitation

See https://github.com/actions/setup-python/issues/544
---
 .github/workflows/unit-tests.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/unit-tests.yml b/.github/workflows/unit-tests.yml
index 97eecdb84..42af21e53 100644
--- a/.github/workflows/unit-tests.yml
+++ b/.github/workflows/unit-tests.yml
@@ -8,7 +8,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: ['3.6', '3.8', '3.10']
+        python-version: ['3.7', '3.8', '3.10']
 
     steps:
       - name: Checkout 🛎️

From 637f209b7a559672c2cecc8714c63330c8751475 Mon Sep 17 00:00:00 2001
From: Philipp Holl <philipp.holl@tum.de>
Date: Tue, 13 Dec 2022 14:28:37 +0100
Subject: [PATCH 21/29] [math] Fix flatten() interactions

* Remove flatten() call by unstack()
* Fix test_flatten
* Fix test_custom_gradient_vector
---
 phi/math/_magic_ops.py                | 12 ++++--------
 tests/commit/math/test__functional.py |  2 +-
 tests/commit/math/test__magic_ops.py  |  3 ++-
 3 files changed, 7 insertions(+), 10 deletions(-)

diff --git a/phi/math/_magic_ops.py b/phi/math/_magic_ops.py
index e11ad02bd..3d9b1da49 100644
--- a/phi/math/_magic_ops.py
+++ b/phi/math/_magic_ops.py
@@ -31,7 +31,7 @@ def unstack(value, dim: DimFilter):
         # Out: (0.0, 0.0, 0.0, 0.0, 0.0)
         ```
     """
-    assert isinstance(value, Sliceable) and isinstance(value, Shaped)
+    assert isinstance(value, Sliceable) and isinstance(value, Shaped), f"Cannot unstack {type(value).__name__}. Must be Sliceable and Shaped, see https://tum-pbs.github.io/PhiFlow/phi/math/magic.html"
     dims = shape(value).only(dim)
     assert dims.rank > 0, "unstack() requires at least one dimension"
     if dims.rank == 1:
@@ -45,9 +45,9 @@ def unstack(value, dim: DimFilter):
     else:  # multiple dimensions
         if hasattr(value, '__pack_dims__'):
             packed_dim = batch('_unstack')
-            value = value.__pack_dims__(dims.names, packed_dim, pos=None)
-            if value is not NotImplemented:
-                return unstack(value, packed_dim)
+            value_packed = value.__pack_dims__(dims.names, packed_dim, pos=None)
+            if value_packed is not NotImplemented:
+                return unstack(value_packed, packed_dim)
         first_unstacked = unstack(value, dims[0])
         inner_unstacked = [unstack(v, dims.without(dims[0])) for v in first_unstacked]
         return sum(inner_unstacked, ())
@@ -342,10 +342,6 @@ def pack_dims(value, dims: DimFilter, packed_dim: Shape, pos: int or None = None
         return value if packed_dim.size is None else expand(value, packed_dim, **kwargs)  # Inserting size=1 can cause shape errors
     elif len(dims) == 1:
         return rename_dims(value, dims, packed_dim, **kwargs)
-    if dims.rank == shape(value).rank and hasattr(value, '__flatten__'):
-        result = value.__flatten__(packed_dim, **kwargs)
-        if result is not NotImplemented:
-            return result
     if hasattr(value, '__pack_dims__'):
         result = value.__pack_dims__(dims.names, packed_dim, pos, **kwargs)
         if result is not NotImplemented:
diff --git a/tests/commit/math/test__functional.py b/tests/commit/math/test__functional.py
index 193b90b68..9e68170ef 100644
--- a/tests/commit/math/test__functional.py
+++ b/tests/commit/math/test__functional.py
@@ -121,7 +121,7 @@ def f(x):
             return x.x[:2]
 
         def grad(_inputs, _y, df):
-            return {'x': math.flatten(math.expand(df * 0, batch(tmp=2)))}
+            return {'x': math.flatten(math.expand(df * 0, batch(tmp=2)), flatten_batch=True)}
 
         def loss(x):
             fg = math.custom_gradient(f, grad)
diff --git a/tests/commit/math/test__magic_ops.py b/tests/commit/math/test__magic_ops.py
index c01863ff5..7421a5aa2 100644
--- a/tests/commit/math/test__magic_ops.py
+++ b/tests/commit/math/test__magic_ops.py
@@ -154,7 +154,8 @@ def test_unpack_dim(self):
     def test_flatten(self):
         for test_class in TEST_CLASSES:
             a = test_class(spatial(x=5) & batch(b=2))
-            self.assertEqual(instance(points=10), flatten(a, instance('points')).shape)
+            self.assertEqual(instance(points=10), flatten(a, instance('points'), flatten_batch=True).shape)
+            self.assertEqual(batch(b=2) & instance(points=5), flatten(a, instance('points'), flatten_batch=False).shape)
 
     def test_bound_dim(self):
         for test_class in TEST_CLASSES:

From 3f7df65c8014fc3611c3d9e31809cee682c02570 Mon Sep 17 00:00:00 2001
From: Philipp Holl <philipp.holl@tum.de>
Date: Tue, 13 Dec 2022 15:50:29 +0100
Subject: [PATCH 22/29] [math] Avoid Tensor-valued sizes in
 Shape.after_gather()

---
 phi/math/_shape.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/phi/math/_shape.py b/phi/math/_shape.py
index 4e83e3f42..52d162909 100644
--- a/phi/math/_shape.py
+++ b/phi/math/_shape.py
@@ -52,6 +52,9 @@ def __init__(self, sizes: tuple, names: tuple, types: tuple, item_names: tuple):
         # assert isinstance(self.item_names, tuple)
         # assert all([items is None or isinstance(items, tuple) for items in self.item_names])
         # assert all([items is None or all([isinstance(n, str) for n in items]) for items in self.item_names])
+        # for size in sizes:
+        #     if size is not None and not isinstance(size, int):
+        #         assert size.rank > 0
 
     def _to_dict(self, include_sizes=True):
         result = dict(names=self.names, types=self.types, item_names=self.item_names)
@@ -955,6 +958,7 @@ def after_gather(self, selection: dict) -> 'Shape':
                 else:
                     from phi.math import Tensor
                     gathered_sizes = [(s[{sel_dim: selection}] if isinstance(s, Tensor) else s) for s in result.sizes]
+                    gathered_sizes = [(int(s) if isinstance(s, Tensor) and s.rank == 0 else s) for s in gathered_sizes]
                     result = result.with_sizes(gathered_sizes, keep_item_names=True).without(sel_dim)
             elif isinstance(selection, slice):
                 step = selection.step or 1

From 8f1db9b3aca3f55411f1b13da125156af8f24008 Mon Sep 17 00:00:00 2001
From: Philipp Holl <philipp.holl@tum.de>
Date: Tue, 13 Dec 2022 16:01:33 +0100
Subject: [PATCH 23/29] [math] Fix constant padding bug

This bug occurred when padding with a vector-valued extrapolation value.
---
 phi/math/_ops.py          |  8 +++++---
 phi/math/extrapolation.py | 10 +++++++---
 2 files changed, 12 insertions(+), 6 deletions(-)

diff --git a/phi/math/_ops.py b/phi/math/_ops.py
index 0c76d7250..bd69c983b 100644
--- a/phi/math/_ops.py
+++ b/phi/math/_ops.py
@@ -740,13 +740,15 @@ def pad(value: Tensor, widths: dict, mode: 'e_.Extrapolation' or Tensor or Numbe
         Padded `Tensor`
     """
     mode = mode if isinstance(mode, e_.Extrapolation) else e_.ConstantExtrapolation(mode)
-    has_negative_widths = any(w[0] < 0 or w[1] < 0 for w in widths.values())
+    has_negative_widths = any(w0 < 0 or w1 < 0 for w0, w1 in widths.values())
+    has_positive_widths = any(w0 > 0 or w1 > 0 for w0, w1 in widths.values())
     slices = None
     if has_negative_widths:
         slices = {dim: slice(max(0, -w[0]), min(0, w[1]) or None) for dim, w in widths.items()}
         widths = {dim: (max(0, w[0]), max(0, w[1])) for dim, w in widths.items()}
-    result = mode.pad(value, widths, **kwargs)
-    return result[slices] if has_negative_widths else result
+    result_padded = mode.pad(value, widths, **kwargs) if has_positive_widths else value
+    result_sliced = result_padded[slices] if has_negative_widths else result_padded
+    return result_sliced
 
 
 def closest_grid_values(grid: Tensor,
diff --git a/phi/math/extrapolation.py b/phi/math/extrapolation.py
index 4e819b32f..09bdcddd1 100644
--- a/phi/math/extrapolation.py
+++ b/phi/math/extrapolation.py
@@ -248,7 +248,11 @@ def pad(self, value: Tensor, widths: dict, **kwargs):
             native = value._native
             ordered_pad_widths = order_by_shape(value.shape, widths, default=(0, 0))
             backend = choose_backend(native, pad_value.native())
-            result_tensor = backend.pad(native, ordered_pad_widths, 'constant', pad_value.native())
+            for dim in pad_value.shape.non_batch.names:
+                assert dim in value.shape, f"Cannot pad tensor {value.shape} with extrapolation {pad_value.shape} because non-batch dimension '{dim}' is missing."
+            result_tensor = NotImplemented
+            if pad_value.rank == 0:
+                result_tensor = backend.pad(native, ordered_pad_widths, 'constant', pad_value.native())
             if result_tensor is NotImplemented:
                 return Extrapolation.pad(self, value, widths, **kwargs)
             return NativeTensor(result_tensor, value.shape.after_pad(widths))
@@ -267,8 +271,8 @@ def pad(self, value: Tensor, widths: dict, **kwargs):
         elif isinstance(value, TensorStack):
             if not value.requires_broadcast:
                 return self.pad(value._cache(), widths)
-            inner_widths = {dim: w for dim, w in widths.items() if dim != value.stack_dim_name}
-            tensors = [self.pad(t, inner_widths) for t in value.dimension(value.stack_dim.name)]
+            inner_widths = {dim: w for dim, w in widths.items() if dim != value.stack_dim.name}
+            tensors = [self[{value.stack_dim.name: i}].pad(t, inner_widths) for i, t in enumerate(value.dimension(value.stack_dim.name))]
             return TensorStack(tensors, value.stack_dim)
         else:
             return Extrapolation.pad(self, value, widths, **kwargs)

From 9395f1745c66d41cb53ae9abd7e180c416dd1967 Mon Sep 17 00:00:00 2001
From: Philipp Holl <philipp.holl@tum.de>
Date: Tue, 13 Dec 2022 19:59:21 +0100
Subject: [PATCH 24/29] [math] Add expand_values option to stack()

vec() now uses this mode by default.

* Add unit tests
---
 phi/math/_magic_ops.py               | 27 +++++++++++++++++++--------
 phi/math/_nd.py                      |  4 ++--
 tests/commit/math/test__magic_ops.py |  9 ++++++---
 tests/commit/math/test__nd.py        |  4 ++++
 4 files changed, 31 insertions(+), 13 deletions(-)

diff --git a/phi/math/_magic_ops.py b/phi/math/_magic_ops.py
index 3d9b1da49..f2e4053b4 100644
--- a/phi/math/_magic_ops.py
+++ b/phi/math/_magic_ops.py
@@ -53,7 +53,7 @@ def unstack(value, dim: DimFilter):
         return sum(inner_unstacked, ())
 
 
-def stack(values: tuple or list or dict, dim: Shape, **kwargs):
+def stack(values: tuple or list or dict, dim: Shape, expand_values=False, **kwargs):
     """
     Stacks `values` along the new dimension `dim`.
     All values must have the same spatial, instance and channel dimensions. If the dimension sizes vary, the resulting tensor will be non-uniform.
@@ -68,6 +68,9 @@ def stack(values: tuple or list or dict, dim: Shape, **kwargs):
         dim: `Shape` with a least one dimension. None of these dimensions can be present with any of the `values`.
             If `dim` is a single-dimension shape, its size is determined from `len(values)` and can be left undefined (`None`).
             If `dim` is a multi-dimension shape, its volume must be equal to `len(values)`.
+        expand_values: If `True`, will first add missing dimensions to all values, not just batch dimensions.
+            This allows tensors with different dimensions to be stacked.
+            The resulting tensor will have all dimensions that are present in `values`.
         **kwargs: Additional keyword arguments required by specific implementations.
             Adding spatial dimensions to fields requires the `bounds: Box` argument specifying the physical extent of the new dimensions.
             Adding batch dimensions must always work without keyword arguments.
@@ -91,14 +94,22 @@ def stack(values: tuple or list or dict, dim: Shape, **kwargs):
     assert len(values) > 0, f"stack() got empty sequence {values}"
     assert isinstance(dim, Shape)
     values_ = tuple(values.values()) if isinstance(values, dict) else values
-    for v in values_[1:]:
-        assert set(non_batch(v).names) == set(non_batch(values_[0]).names), f"Stacked values must have the same non-batch dimensions but got {non_batch(values_[0])} and {non_batch(v)}"
-    # --- Add missing batch dimensions ---
-    all_batch_dims = merge_shapes(*[batch(v) for v in values_])
-    if isinstance(values, dict):
-        values = {k: expand(v, all_batch_dims) for k, v in values.items()}
+    if not expand_values:
+        for v in values_[1:]:
+            assert set(non_batch(v).names) == set(non_batch(values_[0]).names), f"Stacked values must have the same non-batch dimensions but got {non_batch(values_[0])} and {non_batch(v)}"
+    # --- Add missing dimensions ---
+    if expand_values:
+        all_dims = merge_shapes(*values_)
+        if isinstance(values, dict):
+            values = {k: expand(v, all_dims.without(shape(v).non_batch)) for k, v in values.items()}
+        else:
+            values = [expand(v, all_dims.without(shape(v).non_batch)) for v in values]
     else:
-        values = [expand(v, all_batch_dims) for v in values]
+        all_batch_dims = merge_shapes(*[batch(v) for v in values_])
+        if isinstance(values, dict):
+            values = {k: expand(v, all_batch_dims) for k, v in values.items()}
+        else:
+            values = [expand(v, all_batch_dims) for v in values]
     if dim.rank == 1:
         assert dim.size == len(values) or dim.size is None, f"stack dim size must match len(values) or be undefined but got {dim} for {len(values)} values"
         if dim.size is None:
diff --git a/phi/math/_nd.py b/phi/math/_nd.py
index f758095f5..0c4e977d3 100644
--- a/phi/math/_nd.py
+++ b/phi/math/_nd.py
@@ -5,7 +5,7 @@
 from . import _ops as math
 from . import extrapolation as extrapolation
 from ._magic_ops import stack, rename_dims, concat, variable_values
-from ._shape import Shape, channel, batch, spatial, DimFilter, parse_dim_order, shape
+from ._shape import Shape, channel, batch, spatial, DimFilter, parse_dim_order, shape, merge_shapes
 from ._tensors import Tensor, wrap
 from .magic import PhiTreeNode
 from .extrapolation import Extrapolation
@@ -22,7 +22,7 @@ def vec(name='vector', **components) -> Tensor:
     Returns:
         `Tensor`
     """
-    return stack(components, channel(name))
+    return stack(components, channel(name), expand_values=True)
 
 
 def const_vec(value: float or Tensor, dim: Shape or tuple or list or str):
diff --git a/tests/commit/math/test__magic_ops.py b/tests/commit/math/test__magic_ops.py
index 7421a5aa2..1e1e93a4c 100644
--- a/tests/commit/math/test__magic_ops.py
+++ b/tests/commit/math/test__magic_ops.py
@@ -2,7 +2,7 @@
 from unittest import TestCase
 
 from phi.math import batch, unstack, Shape, merge_shapes, stack, concat, expand, spatial, shape, instance, rename_dims, \
-    pack_dims, random_normal, flatten, unpack_dim, EMPTY_SHAPE, Tensor, Dict, channel
+    pack_dims, random_normal, flatten, unpack_dim, EMPTY_SHAPE, Tensor, Dict, channel, linspace
 from phi.math.magic import BoundDim, Shaped, Sliceable, Shapable, PhiTreeNode, slicing_dict
 
 
@@ -111,11 +111,14 @@ def test_unstack(self):
 
     def test_stack(self):
         for test_class in TEST_CLASSES:
-            test_class = TEST_CLASSES[1]
             a = test_class(spatial(x=5))
-            # self.assertEqual(spatial(x=5) & batch(b=2), stack([a, a], batch('b')).shape)
+            self.assertEqual(spatial(x=5) & batch(b=2), stack([a, a], batch('b')).shape)
             self.assertEqual(spatial(x=5) & batch(b='a1,a2'), stack({'a1': a, 'a2': a}, batch('b')).shape)
 
+    def test_stack_expand(self):
+        v = stack([0, linspace(0, 1, instance(points=10))], channel(vector='x,y'), expand_values=True)
+        self.assertEqual(instance(points=10) & channel(vector='x,y'), v.shape)
+
     def test_multi_dim_stack(self):
         for test_class in TEST_CLASSES:
             a = test_class(spatial(x=5))
diff --git a/tests/commit/math/test__nd.py b/tests/commit/math/test__nd.py
index 2c80e2a97..21065551e 100644
--- a/tests/commit/math/test__nd.py
+++ b/tests/commit/math/test__nd.py
@@ -255,3 +255,7 @@ def test_vector_length(self):
     def test_dim_mask(self):
         math.assert_close((1, 0, 0), math.dim_mask(spatial('x,y,z'), 'x'))
         math.assert_close((1, 0, 1), math.dim_mask(spatial('x,y,z'), 'x,z'))
+
+    def test_vec_expand(self):
+        v = math.vec(x=0, y=math.linspace(0, 1, instance(points=10)))
+        self.assertEqual(instance(points=10) & channel(vector='x,y'), v.shape)

From e680eac890ae367967475675e6293de040f93f7b Mon Sep 17 00:00:00 2001
From: Philipp Holl <philipp.holl@tum.de>
Date: Tue, 13 Dec 2022 19:59:51 +0100
Subject: [PATCH 25/29] =?UTF-8?q?[=CE=A6]=20Bump=20version=20to=202.2.6?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 phi/VERSION | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/phi/VERSION b/phi/VERSION
index ecf00d901..1506473e2 100644
--- a/phi/VERSION
+++ b/phi/VERSION
@@ -1 +1 @@
-2.2.5
\ No newline at end of file
+2.2.6
\ No newline at end of file

From 58e278acc91e5b1646eb06eebad3734ca84ca662 Mon Sep 17 00:00:00 2001
From: Philipp Holl <philipp.holl@tum.de>
Date: Tue, 13 Dec 2022 21:15:42 +0100
Subject: [PATCH 26/29] [vis] Fix matplotlib 3D point cloud plot

---
 phi/vis/_matplotlib/_matplotlib_plots.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/phi/vis/_matplotlib/_matplotlib_plots.py b/phi/vis/_matplotlib/_matplotlib_plots.py
index ef3262550..9724be9da 100644
--- a/phi/vis/_matplotlib/_matplotlib_plots.py
+++ b/phi/vis/_matplotlib/_matplotlib_plots.py
@@ -248,7 +248,7 @@ def _plot(axis, data: SampledField, space: Box, show_color_bar, vmin, vmax, **pl
         if data.points.shape.non_channel.rank > 1:
             data_list = field.unstack(data, data.points.shape.non_channel[0].name)
             for d in data_list:
-                _plot(axis, d, show_color_bar, vmin, vmax, **plt_args)
+                _plot(axis, d, space, show_color_bar, vmin, vmax, **plt_args)
         else:
             x, y, z = math.reshaped_numpy(data.points.vector[dims], [vector, data.shape.non_channel])
             color = [d.native() for d in data.color.points.unstack(len(x))]

From 75741d1488a2150e46ebbf15751637dfb617df2a Mon Sep 17 00:00:00 2001
From: Philipp Holl <philipp.holl@tum.de>
Date: Wed, 14 Dec 2022 13:09:47 +0100
Subject: [PATCH 27/29] [doc] Add examples math docstrings

vec(), tensor(), linspace(), pad()
---
 phi/math/_nd.py      | 12 ++++++++++++
 phi/math/_ops.py     | 18 ++++++++++++++++++
 phi/math/_tensors.py | 32 +++++++++++++++++++++++++-------
 3 files changed, 55 insertions(+), 7 deletions(-)

diff --git a/phi/math/_nd.py b/phi/math/_nd.py
index 0c4e977d3..f83d19651 100644
--- a/phi/math/_nd.py
+++ b/phi/math/_nd.py
@@ -21,6 +21,18 @@ def vec(name='vector', **components) -> Tensor:
 
     Returns:
         `Tensor`
+
+    Examples:
+        ```python
+        vec(x=1, y=0, z=-1)
+        # Out: (x=1, y=0, z=-1)
+
+        vec(x=1., z=0)
+        # Out: (x=1.000, z=0.000)
+
+        vec(x=tensor([1, 2, 3], instance('particles')), y=0)
+        # Out: (x=1, y=0); (x=2, y=0); (x=3, y=0) (particlesⁱ=3, vectorᶜ=x,y)
+        ```
     """
     return stack(components, channel(name), expand_values=True)
 
diff --git a/phi/math/_ops.py b/phi/math/_ops.py
index bd69c983b..3238daa9c 100644
--- a/phi/math/_ops.py
+++ b/phi/math/_ops.py
@@ -630,6 +630,15 @@ def linspace(start: int or Tensor, stop, dim: Shape) -> Tensor:
 
     Returns:
         `Tensor`
+
+    Examples:
+        ```python
+        math.linspace(0, 1, spatial(x=5))
+        # Out: (0.000, 0.250, 0.500, 0.750, 1.000) along xˢ
+
+        math.linspace(0, (-1, 1), spatial(x=3))
+        # Out: (0.000, 0.000); (-0.500, 0.500); (-1.000, 1.000) (xˢ=3, vectorᶜ=2)
+        ```
     """
     assert isinstance(dim, Shape) and dim.rank == 1, f"dim must be a single-dimension Shape but got {dim}"
     if is_scalar(start) and is_scalar(stop):
@@ -738,6 +747,15 @@ def pad(value: Tensor, widths: dict, mode: 'e_.Extrapolation' or Tensor or Numbe
 
     Returns:
         Padded `Tensor`
+
+    Examples:
+        ```python
+        math.pad(math.ones(spatial(x=10, y=10)), {'x': (1, 1), 'y': (2, 1)}, 0)
+        # Out: (xˢ=12, yˢ=13) 0.641 ± 0.480 (0e+00...1e+00)
+
+        math.pad(math.ones(spatial(x=10, y=10)), {'x': (1, -1)}, 0)
+        # Out: (xˢ=10, yˢ=10) 0.900 ± 0.300 (0e+00...1e+00)
+        ```
     """
     mode = mode if isinstance(mode, e_.Extrapolation) else e_.ConstantExtrapolation(mode)
     has_negative_widths = any(w0 < 0 or w1 < 0 for w0, w1 in widths.values())
diff --git a/phi/math/_tensors.py b/phi/math/_tensors.py
index d247c1a74..63987b9dd 100644
--- a/phi/math/_tensors.py
+++ b/phi/math/_tensors.py
@@ -1599,17 +1599,35 @@ def tensor(data: Tensor or Shape or tuple or list or numbers.Number,
         `phi.math.wrap()` which uses `convert=False`, `layout()`.
 
     Args:
-      data: native tensor, scalar, sequence, Shape or Tensor
-      shape: Ordered dimensions and types. If sizes are defined, they will be checked against `data`.`
-      convert: If True, converts the data to the native format of the current default backend.
-        If False, wraps the data in a `Tensor` but keeps the given data reference if possible.
+        data: native tensor, scalar, sequence, Shape or Tensor
+        shape: Ordered dimensions and types. If sizes are defined, they will be checked against `data`.`
+        convert: If True, converts the data to the native format of the current default backend.
+            If False, wraps the data in a `Tensor` but keeps the given data reference if possible.
 
     Raises:
-      AssertionError: if dimension names are not provided and cannot automatically be inferred
-      ValueError: if `data` is not tensor-like
+        AssertionError: if dimension names are not provided and cannot automatically be inferred
+        ValueError: if `data` is not tensor-like
 
     Returns:
-      Tensor containing same values as data
+        Tensor containing same values as data
+
+    Examples:
+        ```python
+        tensor([1, 2, 3], channel(vector='x,y,z'))
+        # Out: (x=1, y=2, z=3)
+
+        tensor([1., 2, 3], channel(vector='x,y,z'))
+        # Out: (x=1.000, y=2.000, z=3.000) float64
+
+        tensor(numpy.zeros([10, 8, 6, 2]), batch('batch'), spatial('x,y'), channel(vector='x,y'))
+        # Out: (batchᵇ=10, xˢ=8, yˢ=6, vectorᶜ=x,y) float64 const 0.0
+
+        tensor([(0, 1), (0, 2), (1, 3)], instance('particles'), channel(vector='x,y'))
+        # Out: (x=0, y=1); (x=0, y=2); (x=1, y=3) (particlesⁱ=3, vectorᶜ=x,y)
+
+        tensor(numpy.random.randn(10))
+        # Out: (vectorᶜ=10) float64 -0.128 ± 1.197 (-2e+00...2e+00)
+        ```
     """
     assert all(isinstance(s, Shape) for s in shape), f"Cannot create tensor because shape needs to be one or multiple Shape instances but got {shape}"
     shape = None if len(shape) == 0 else concat_shapes(*shape)

From c24e82164cd56b07cc2922522998ceef189de736 Mon Sep 17 00:00:00 2001
From: Philipp Holl <philipp.holl@tum.de>
Date: Wed, 14 Dec 2022 23:28:06 +0100
Subject: [PATCH 28/29] [math] Improved support for Tensor-tuple operations

Operators are now supported between Tensors and tuples containing Tensors and natives/constants.

* Add unit tests
---
 phi/math/_tensors.py               | 20 +++++++++++++++++---
 tests/commit/math/test__tensors.py | 12 ++++++++++++
 2 files changed, 29 insertions(+), 3 deletions(-)

diff --git a/phi/math/_tensors.py b/phi/math/_tensors.py
index 63987b9dd..23061e280 100644
--- a/phi/math/_tensors.py
+++ b/phi/math/_tensors.py
@@ -9,7 +9,7 @@
 import numpy
 import numpy as np
 
-from ._magic_ops import PhiTreeNodeType, variable_attributes, copy_with
+from ._magic_ops import PhiTreeNodeType, variable_attributes, copy_with, stack
 from ._shape import (Shape,
                      CHANNEL_DIM, BATCH_DIM, SPATIAL_DIM, EMPTY_SHAPE,
                      parse_dim_order, shape_stack, merge_shapes, channel, concat_shapes,
@@ -671,7 +671,19 @@ def __iter__(self):
     def _tensor(self, other):
         if isinstance(other, Tensor):
             return other
-        return compatible_tensor(other, compat_shape=self.shape, compat_natives=self._natives(), convert=False)
+        elif isinstance(other, (tuple, list)) and any(isinstance(v, Tensor) for v in other):
+            if 'vector' in self.shape:
+                outer_dim = self.shape['vector']
+            elif self.shape.channel_rank == 1:
+                outer_dim = self.shape.channel
+            else:
+                raise ValueError(f"Cannot combine tensor of shape {self.shape} with tuple {tuple([type(v).__name__ for v in other])}")
+            remaining_shape = self.shape.without(outer_dim)
+            other_items = [v if isinstance(v, Tensor) else compatible_tensor(v, compat_shape=remaining_shape, compat_natives=self._natives(), convert=False) for v in other]
+            other_stacked = stack(other_items, outer_dim, expand_values=True)
+            return other_stacked
+        else:
+            return compatible_tensor(other, compat_shape=self.shape, compat_natives=self._natives(), convert=False)
 
     def _op1(self, native_function):
         """
@@ -1324,6 +1336,8 @@ def _op2(self, other, operator, native_function, op_name: str = 'unknown', op_sy
             else:
                 combined_shape = (self._shape & other_t._shape).with_sizes(inner.shape)
                 return CollapsedTensor(inner, combined_shape)
+        elif not isinstance(other, Tensor):  # was converted to Tensor, probably TensorStack
+            return operator(self, other_t)
         else:
             return NotImplemented
 
@@ -1798,7 +1812,7 @@ def compatible_tensor(data, compat_shape: Shape = None, compat_natives=(), conve
             warnings.warn(f"Combining a phi.math.Tensor with a {data_type} of same shape is not invariant under shape permutations. Please convert the {data_type} to a phi.math.Tensor first. Shapes: {shape} and {compat_shape}", SyntaxWarning, stacklevel=5)
             return NativeTensor(data, compat_shape.with_sizes(shape))
         else:
-            raise ValueError(f"Cannot combine native tensor of shape {shape} with tensor of shape {compat_shape}")
+            raise ValueError(f"Cannot combine tensor of shape {shape} with tensor of shape {compat_shape}")
 
 
 def broadcastable_native_tensors(*tensors):
diff --git a/tests/commit/math/test__tensors.py b/tests/commit/math/test__tensors.py
index 0533a375e..d4f73a7d5 100644
--- a/tests/commit/math/test__tensors.py
+++ b/tests/commit/math/test__tensors.py
@@ -605,3 +605,15 @@ def test_numpy_function_interface(self):
                 self.assertFalse(np.equal(x, y))
                 self.assertTrue(x != y)
                 self.assertTrue(np.not_equal(x, y))
+
+    def test_broadcast_vector_tuple(self):
+        t = vec('something', x=0, y=1)
+        i = math.linspace(0, 1, instance(particles=10))
+        result = t + (0, i)
+        self.assertEqual(instance(particles=10) & channel(something='x,y'), result.shape)
+
+    def test_broadcast_matrix_tuple(self):
+        i = math.linspace(0, 1, instance(particles=10))
+        t = math.expand(vec(x=0, y=1), channel(c=2))
+        result = t + (0, i)
+        self.assertEqual(instance(particles=10) & channel(vector='x,y', c=2), result.shape)

From f6520e44c300f2a1aa44aa11ab5d13e80184b7ed Mon Sep 17 00:00:00 2001
From: Philipp Holl <philipp.holl@tum.de>
Date: Thu, 15 Dec 2022 19:13:05 +0100
Subject: [PATCH 29/29] [doc] Update publications & documentation links in
 README.md

---
 README.md | 16 ++++++++++++----
 1 file changed, 12 insertions(+), 4 deletions(-)

diff --git a/README.md b/README.md
index 10bdc959d..795c81964 100644
--- a/README.md
+++ b/README.md
@@ -50,9 +50,10 @@ See the [detailed installation instructions](https://tum-pbs.github.io/PhiFlow/I
 
 To get started, check out our YouTube tutorial series and the following Jupyter notebooks:
 
-* [<img src="https://www.tensorflow.org/images/colab_logo_32px.png" height=16> Fluids](https://colab.research.google.com/github/tum-pbs/PhiFlow/blob/develop/docs/Fluids_Tutorial.ipynb): Introduction to core classes and fluid-related functions.
-* [<img src="https://www.tensorflow.org/images/colab_logo_32px.png" height=16> Solar System](https://colab.research.google.com/github/tum-pbs/PhiFlow/blob/develop/docs/Planets_Tutorial.ipynb): Visualize a many-body system with Newtonian gravity.
-* [<img src="https://www.tensorflow.org/images/colab_logo_32px.png" height=16> Learning to Throw](https://colab.research.google.com/github/tum-pbs/PhiFlow/blob/develop/docs/Learn_to_Throw_Tutorial.ipynb): Train a neural network to hit a target, comparing supervised and differentiable physics losses.
+* [<img src="https://www.tensorflow.org/images/colab_logo_32px.png" height=16>](https://colab.research.google.com/github/tum-pbs/PhiFlow/blob/develop/docs/Math_Introduction.ipynb) [Tensors](https://tum-pbs.github.io/PhiFlow/Math_Introduction.html): Introduction to tensors.
+* [<img src="https://www.tensorflow.org/images/colab_logo_32px.png" height=16>](https://colab.research.google.com/github/tum-pbs/PhiFlow/blob/develop/docs/Fluids_Tutorial.ipynb) [Fluids](https://tum-pbs.github.io/PhiFlow/Fluids_Tutorial.html): Introduction to core classes and fluid-related functions.
+* [<img src="https://www.tensorflow.org/images/colab_logo_32px.png" height=16>](https://colab.research.google.com/github/tum-pbs/PhiFlow/blob/develop/docs/Planets_Tutorial.ipynb) [Solar System](https://tum-pbs.github.io/PhiFlow/Planets_Tutorial.html): Visualize a many-body system with Newtonian gravity.
+* [<img src="https://www.tensorflow.org/images/colab_logo_32px.png" height=16>](https://colab.research.google.com/github/tum-pbs/PhiFlow/blob/develop/docs/Learn_to_Throw_Tutorial.ipynb) [Learn to Throw](https://tum-pbs.github.io/PhiFlow/Learn_to_Throw_Tutorial.html): Train a neural network to hit a target, comparing supervised and differentiable physics losses.
 
 If you like to work with an IDE, like PyCharm or VS Code, the following demos will also be helpful:
 
@@ -61,13 +62,20 @@ If you like to work with an IDE, like PyCharm or VS Code, the following demos wi
 
 ## Publications
 
-We have recently submitted a whitepaper.
+We will upload a whitepaper, soon.
 In the meantime, please cite the ICLR 2020 paper.
 
 * [Learning to Control PDEs with Differentiable Physics](https://ge.in.tum.de/publications/2020-iclr-holl/), *Philipp Holl, Vladlen Koltun, Nils Thuerey*, ICLR 2020.
 * [Solver-in-the-Loop: Learning from Differentiable Physics to Interact with Iterative PDE-Solvers](https://arxiv.org/abs/2007.00016), *Kiwon Um, Raymond Fei, Philipp Holl, Robert Brand, Nils Thuerey*, NeurIPS 2020.
 * [Φ<sub>Flow</sub>: A Differentiable PDE Solving Framework for Deep Learning via Physical Simulations](https://montrealrobotics.ca/diffcvgp/), *Nils Thuerey, Kiwon Um, Philipp Holl*, DiffCVGP workshop at NeurIPS 2020.
+* [Physics-based Deep Learning](https://physicsbaseddeeplearning.org/intro.html) (book), *Nils Thuerey, Philipp Holl, Maximilian Mueller, Patrick Schnell, Felix Trost, Kiwon Um*, DiffCVGP workshop at NeurIPS 2020.
 * [Half-Inverse Gradients for Physical Deep Learning](https://arxiv.org/abs/2203.10131), *Patrick Schnell, Philipp Holl, Nils Thuerey*, ICLR 2022.
+* [Scale-invariant Learning by Physics Inversion](https://arxiv.org/abs/2109.15048), *Philipp Holl, Vladlen Koltun, Nils Thuerey*, NeurIPS 2022.
+
+Φ<sub>Flow</sub> has been used in the following data sets:
+
+* [PDEBench](https://github.com/pdebench/PDEBench)
+* [PDEarena](https://microsoft.github.io/pdearena/)
 
 ## Version History