Issue with fluid.make_incompressible in Navier-Stokes Forward Simulation

[projectavi]

def step(velocity, smoke, pressure, dt=1.0, buoyancy_factor=1.0):
    smoke = advect.semi_lagrangian(smoke, velocity, dt) + INFLOW
    buoyancy_force = (smoke * (0, buoyancy_factor)).at(velocity) # Resample the smoke to the points where velocity was sampled
    velocity = advect.semi_lagrangian(velocity, velocity, dt) + dt * buoyancy_force
    velocity = diffuse.explicit(velocity, NU, dt)
    velocity, pressure = fluid.make_incompressible(velocity)
    return velocity, smoke, pressure

velocity, smoke, pressure = step(velocity, smoke, None, dt=DT)

print("Max. velocity and mean marker density: " + format( [ math.max(velocity.values) , math.mean(smoke.values) ] ))

pylab.imshow(np.asarray(smoke.values.numpy('y,x')), origin='lower', cmap='magma')

In this code from the Navier-Stokes forward simulation document I am getting an error raised from the line velocity, pressure = fluid.make_incompressible(velocity)

SparseEfficiencyWarning: spsolve requires A be CSC or CSR matrix format x = spsolve(lin[batch], y[batch]) # returns nan when diverges

Followed by the full error message.

---------------------------------------------------------------------------
Diverged                                  Traceback (most recent call last)
Cell In[5], line 11
      8     velocity, pressure = fluid.make_incompressible(velocity)
      9     return velocity, smoke, pressure
---> 11 velocity, smoke, pressure = step(velocity, smoke, None, dt=DT)
     13 print("Max. velocity and mean marker density: " + format( [ math.max(velocity.values) , math.mean(smoke.values) ] ))
     15 pylab.imshow(np.asarray(smoke.values.numpy('y,x')), origin='lower', cmap='magma')

Cell In[5], line 8, in step(velocity, smoke, pressure, dt, buoyancy_factor)
      6 velocity = advect.semi_lagrangian(velocity, velocity, dt) + dt * buoyancy_force
      7 velocity = diffuse.explicit(velocity, NU, dt)
----> 8 velocity, pressure = fluid.make_incompressible(velocity)
      9 return velocity, smoke, pressure

File ~/miniconda3/envs/PBDL/lib/python3.12/site-packages/phi/physics/fluid.py:156, in make_incompressible(velocity, obstacles, solve, active, order, correct_skew, wide_stencil)
    154 if wide_stencil is None:
    155     wide_stencil = not velocity.is_staggered
--> 156 pressure = math.solve_linear(masked_laplace, div, solve, velocity.boundary, hard_bcs, active, wide_stencil=wide_stencil, order=order, implicit=None, upwind=None, correct_skew=correct_skew)
    157 # --- Subtract grad p ---
    158 grad_pressure = field.spatial_gradient(pressure, input_velocity.extrapolation, at=velocity.sampled_at, order=order, scheme='green-gauss')

File ~/miniconda3/envs/PBDL/lib/python3.12/site-packages/phiml/math/_optimize.py:671, in solve_linear(f, y, solve, grad_for_f, f_kwargs, *f_args, **f_kwargs_)
    668         return result  # must return exactly `x` so gradient isn't computed w.r.t. other quantities
    670     _matrix_solve = attach_gradient_solve(_matrix_solve_forward, auxiliary_args=f'is_backprop,solve{",matrix" if matrix.backend == NUMPY else ""}', matrix_adjoint=grad_for_f)
--> 671     return _matrix_solve(y - bias, solve, matrix)
    672 else:  # Matrix-free solve
    673     f_args = cached(f_args)

File ~/miniconda3/envs/PBDL/lib/python3.12/site-packages/phiml/math/_functional.py:952, in CustomGradientFunction.__call__(self, *args, **kwargs)
    950 key, _, natives, _, _ = key_from_args(args, kwargs, self.f_params, cache=False, aux=self.auxiliary_args, attr_type=variable_attributes)
    951 if not key.backend.supports(Backend.jacobian) and not key.backend.supports(Backend.jacobian):
--> 952     return self.f(*args, **kwargs)  # no need to use custom gradient if gradients aren't supported anyway
    953 elif not key.backend.supports(Backend.custom_gradient):
    954     warnings.warn(f"custom_gradient() not supported by {key.backend}. Running function '{f_name(self.f)}' as-is.", RuntimeWarning)

File ~/miniconda3/envs/PBDL/lib/python3.12/site-packages/phiml/math/_optimize.py:667, in solve_linear.<locals>._matrix_solve_forward(y, solve, matrix, is_backprop)
    665     idx = b.concat([idx, new_col, new_row], 0)
    666     nat_matrix = b.sparse_coo_tensor(idx, data, (N+1, N+1))
--> 667 result = _linear_solve_forward(y, solve, nat_matrix, pattern_dims_in, pattern_dims_out, preconditioner, backend, is_backprop)
    668 return result

File ~/miniconda3/envs/PBDL/lib/python3.12/site-packages/phiml/math/_optimize.py:780, in _linear_solve_forward(y, solve, native_lin_op, pattern_dims_in, pattern_dims_out, preconditioner, backend, is_backprop)
    778 for tape in _SOLVE_TAPES:
    779     tape._add(solve, trj, result)
--> 780 result.convergence_check(is_backprop and 'TensorFlow' in backend.name)  # raises ConvergenceException
    781 return final_x

File ~/miniconda3/envs/PBDL/lib/python3.12/site-packages/phiml/math/_optimize.py:202, in SolveInfo.convergence_check(self, only_warn)
    200             warnings.warn(self.msg, ConvergenceWarning)
    201         else:
--> 202             raise Diverged(self)
    203 if not self.converged.trajectory[-1].all:
    204     if NotConverged not in self.solve.suppress:

Diverged: Direct solution does not satisfy tolerance: norm(residual)=2.2658770831185393e-05

I am using Python 3.12 and Phiflow 3.4. If I use 3.1 there is a different issue.