jit_compile pytorch does not work with values requiring gradients

[KarlisFre]

Hi, I am trying to do gradient based training, unfortunately with jit_compile I get an error:

RuntimeError: Cannot insert a Tensor that requires grad as a constant. Consider making it a parameter or input, or detaching the gradient

from phi.torch.flow import *

object_geometry = Box(x=(40,60), y=(40, 60), z=(40, 60))
rotation_angle = torch.tensor((0.0, 0.1, 0.1),requires_grad=True)
object_geometry = object_geometry.rotated(rotation_angle)
OBSTACLE = Obstacle(object_geometry, velocity=(0.,0.,0.), angular_velocity=(0.,0.2, 0.2))

velocity = StaggeredGrid((0, 0, 0), 0, x=6, y=6, z=6, bounds=Box(x=100, y=100, z=100))
smoke = CenteredGrid(0, 0, x=6, y=6, z=6, bounds=Box(x=100, y=100, z=100))
INFLOW = 0.2 * resample(Sphere(x=50, y=50, z=10, radius=5), to=smoke, soft=True)
velocity, pressure = fluid.make_incompressible(velocity)

obst_mask = resample(OBSTACLE.geometry, smoke)
plot({"3D": obst_mask})
vis.show()

@jit_compile  # Only for PyTorch, TensorFlow and Jax
def step(v, s, p,obstacle, dt=1.):
    s = advect.mac_cormack(s, v, dt) + INFLOW
    buoyancy = resample(s * (0, 0, 0.1), to=v)
    v = advect.semi_lagrangian(v, v, dt) + buoyancy * dt
    v, p = fluid.make_incompressible(v, (obstacle))
    return v, s, p, obstacle


#for _ in view(smoke, velocity, 'pressure',obst_mask, play=True, namespace=globals(),port=6006).range(warmup=1):
for i in range(10):
    velocity, smoke, pressure, OBSTACLE = step(velocity, smoke, pressure, OBSTACLE)
    obst_mask = resample(OBSTACLE.geometry, smoke)
    print(i)

[holl-]

This should work with Jax and TensorFlow.

JIT-compilation in PyTorch has many problems and pitfalls and is not actually that much faster. For fast code, I recommend jit-compiling with Jax or TensorFlow.

[KarlisFre]

I was considering integrating phiflow with nimblephysics for solid motion and collision support https://github.com/keenon/nimblephysics. But nimblephysics supports only pytorch.

[holl-]

I'm looking into fixing the PyTorch jit but this seems like a tricky problem. I'll get back to you.