So I have a toy attractor model:
points = torch.rand((2**16,3), dtype=torch.float32, device=cuda, requires_grad=False)

def attractor_func(v, a = 10.0, b = 28.0, c = 8.0 / 3.0):
	dv = torch.zeros(points.shape, dtype=torch.float32, device=cuda, requires_grad=False)

	#lorenz
	dv[:, 0] = a * (v[:, 1] - v[:, 0])
	dv[:, 1] = v[:, 0] * (b - v[:, 2]) - v[:, 1]
	dv[:, 2] = v[:, 0] * v[:, 1] - c * v[:, 2]

	return dv

This function is called 4 times from an Runge-Kutta function that is called as quickly as the cpu and gpu can manage.
def rk4(func, h, v):
	k1 = func(v)
	k2 = func(v + (h / 2.0) * k1)
	k3 = func(v + (h / 2.0) * k2)
	k4 = func(v + h * k3)
	return h * (k1 / 6.0 + k2 / 3.0 + k3 / 3.0 + k4 / 6.0)

with this called in between “rendering”:
for _ in range(32):
		points += rk4(attractor_func, 0.01, points)

But something tells me that this is not the most efficient code when it comes to managing GPU memory. I can’t even convert attractor_func to a JIT script, if that is of any use.
Perhaps there is a way to vectorize the dv slices across the second dimension as well?