[REQ] Add adjoints for `texture_sample`

[eigenvivek]

I am trying to differentiable a dense voxelgrid w.r.t. the sampling position. In the example below, the gradient from warp is zero.

import numpy as np
import warp as wp

wp.init()


@wp.kernel
def sample(tex: wp.Texture1D, pos: wp.array(dtype=float), out: wp.array(dtype=float)):
    tid = wp.tid()
    out[tid] = wp.texture_sample(tex, pos[tid], dtype=float)


data = np.arange(16, dtype=np.float32)
tex = wp.Texture1D(data, normalized_coords=False)
pos = wp.array([4.0], dtype=float, requires_grad=True)
out = wp.zeros(1, dtype=float, requires_grad=True)

tape = wp.Tape()
with tape:
    wp.launch(sample, dim=1, inputs=[tex, pos], outputs=[out])
out.grad = wp.ones(1, dtype=float)
tape.backward()

print("out:", out.numpy())            # [3.5]
print("pos.grad:", pos.grad.numpy())  # [0.] expected [1.]

The pytorch version returns the correct gradient (1).

import torch
import torch.nn.functional as F

data = torch.arange(16, dtype=torch.float32).reshape(1, 1, 1, 1, 16)
pos = torch.tensor([4.0], requires_grad=True)

pos_norm = 2.0 * pos / 16.0 - 1.0
grid = torch.stack([pos_norm, torch.zeros_like(pos_norm), torch.zeros_like(pos_norm)], dim=-1)
grid = grid.reshape(1, 1, 1, 1, 3)  # [N, D, H, W, 3]
out = F.grid_sample(data, grid, mode="bilinear", align_corners=False)
out.backward()

print("out:", out.item())             # 3.5
print("pos.grad:", pos.grad.item())   # 1.0

Is texture sampling not currently differentiable or did I implement the gradient incorrectly? Also using the latest nightly.

[shi-eric]

A quick look at the source code and docs indicates that texture_sample is not currently differentiable. For example, the built-in https://nvidia.github.io/warp/language_reference/_generated/warp.texture_sample.html#warp._src.lang.texture_sample has no "Differentiable" tag (just "Kernel"), and the adjoints are currently empty stubs

warp/warp/native/texture.h

Lines 802 to 846 in ccb06a9

	// Adjoint stubs for texture sampling (non-differentiable for now)
	template <typename T>
	CUDA_CALLABLE void
	adj_texture_sample(const texture1d_t& tex, float u, texture1d_t& adj_tex, float& adj_u, const T& adj_ret)
	{
	// Texture sampling is not differentiable in this implementation
	}
	template <typename T>
	CUDA_CALLABLE void
	adj_texture_sample(const texture2d_t& tex, const vec2f& uv, texture2d_t& adj_tex, vec2f& adj_uv, const T& adj_ret)
	{
	// Texture sampling is not differentiable in this implementation
	}
	template <typename T>
	CUDA_CALLABLE void adj_texture_sample(
	const texture2d_t& tex, float u, float v, texture2d_t& adj_tex, float& adj_u, float& adj_v, const T& adj_ret
	)
	{
	// Texture sampling is not differentiable in this implementation
	}
	template <typename T>
	CUDA_CALLABLE void
	adj_texture_sample(const texture3d_t& tex, const vec3f& uvw, texture3d_t& adj_tex, vec3f& adj_uvw, const T& adj_ret)
	{
	// Texture sampling is not differentiable in this implementation
	}
	template <typename T>
	CUDA_CALLABLE void adj_texture_sample(
	const texture3d_t& tex,
	float u,
	float v,
	float w,
	texture3d_t& adj_tex,
	float& adj_u,
	float& adj_v,
	float& adj_w,
	const T& adj_ret
	)
	{
	// Texture sampling is not differentiable in this implementation
	}

We will consider the feature request to support the gradient propagation through the sampling positions.

@nvtw @nvlukasz for viz

[eigenvivek]

Ah I see, thank you for the pointer to the source! An adjoint for textures would be very helpful for a lot of optimization tasks. Will convert this to a request.