dr.norm / dr.squared_norm with axis on Tensor type raises TypeError

[zihay]

Hi,

I’m running into an issue when using dr.norm with a Tensor type. Calling dr.norm (and similarly dr.squared_norm) with an axis argument on Tensorf raises a TypeError, even though the Python stub for these functions exposes an axis parameter. Other reduction operations like dr.sum do work correctly with axis.

Minimal reproducer

import drjit as dr
from drjit.cuda.ad import TensorXf

tensor = dr.ones(TensorXf, shape=(100, 100))
norm = dr.norm(tensor, axis=1)

This currently fails with a TypeError.

I can work around this by manually implementing the norm using dr.sum, but it would be very convenient to have built-in support for axis in dr.norm / dr.squared_norm for tensor types if possible 🙂

[njroussel]

Hi @zihay

Indeeed, the online API reference documentation doesn't mention it, but the stubs incorrectly do have an axis param.
We should support that functionality.