Update phiml imports in phi.flow

phi/flow.py

@@ -12,32 +12,34 @@
 # Modules
 import numpy
 import numpy as np
+import phiml
+from phiml import math, backend
+from phiml.math import extrapolation
 import phi
-from . import math, geom, field, physics, vis
-from .math import extrapolation, backend
+from . import geom, field, physics, vis
 from .physics import fluid, advect, diffuse
 
 # Classes
-from .math import Shape, Tensor, DType, Solve
+from phiml.math import Shape, Tensor, DType, Solve
 from .geom import Geometry, Sphere, Box, Cuboid
 from .field import Field, Grid, CenteredGrid, StaggeredGrid, mask, Noise, PointCloud, Scene, resample, GeometryMask, SoftGeometryMask, HardGeometryMask
 from .vis import Viewer
 from .physics.fluid import Obstacle
 
 # Constants
-from .math import PI, INF, NAN, f
-from .math.extrapolation import PERIODIC, ZERO_GRADIENT
+from phiml.math import PI, INF, NAN, f
+from phiml.math.extrapolation import PERIODIC, ZERO_GRADIENT
 
 # Functions
-from .math import (
+from phiml.math import (
     wrap, tensor, vec,  # Tensor creation
     shape, spatial, channel, batch, instance, dual,
     non_spatial, non_channel, non_batch, non_instance, non_dual,  # Shape functions (magic)
     unstack, stack, concat, expand, rename_dims, pack_dims, unpack_dim, flatten, cast,  # Magic Ops
-    jit_compile, jit_compile_linear, minimize, functional_gradient, solve_linear, solve_nonlinear, iterate, identity,  # jacobian, hessian, custom_gradient # Functional magic
+    jit_compile, jit_compile_linear, minimize, gradient as functional_gradient, gradient, solve_linear, solve_nonlinear, iterate, identity,  # jacobian, hessian, custom_gradient # Functional magic
 )
 from .geom import union
 from .vis import show, view, control, plot
 
 # Exceptions
-from .math import ConvergenceException, NotConverged, Diverged
+from phiml.math import ConvergenceException, NotConverged, Diverged