[geom] Fix Point

phi/geom/_geom.py

@@ -4,7 +4,7 @@
 
 from phi import math
 from phi.math import Tensor, Shape, EMPTY_SHAPE, non_channel, wrap, shape, Extrapolation
-from phiml.math._magic_ops import variable_attributes, expand
+from phiml.math._magic_ops import variable_attributes, expand, stack
 from phi.math.magic import BoundDim, slicing_dict
 
 
@@ -20,6 +20,9 @@ class Geometry:
     All geometry objects support batching.
     Thereby any parameter defining the geometry can be varied along arbitrary batch dims.
     All batch dimensions are listed in Geometry.shape.
+
+    Property getters (`@property`, such as `shape`), save for getters, must not depend on any variables marked as *variable* via `__variable_attrs__()` as these may be `None` during tracing.
+    Equality checks must also take this into account.
     """
 
     @property
@@ -636,14 +639,15 @@ def __init__(self, location: math.Tensor):
         assert 'vector' in location.shape, "location must have a vector dimension"
         assert location.shape.get_item_names('vector') is not None, "Vector dimension needs to list spatial dimension as item names."
         self._location = location
+        self._shape = self._location.shape
 
     @property
     def center(self) -> Tensor:
         return self._location
 
     @property
     def shape(self) -> Shape:
-        return self._location.shape
+        return self._shape
 
     @property
     def faces(self) -> 'Geometry':
@@ -676,6 +680,12 @@ def __hash__(self):
     def __variable_attrs__(self):
         return '_location',
 
+    def __with_attrs__(self, **updates):
+        if '_location' in updates:
+            return Point(updates['_location'])
+        else:
+            return self
+
     @property
     def volume(self) -> Tensor:
         return math.wrap(0)