Compute Norm of Vec3Array, clipped to epsilon.
(self, epsilon: float = 1e-6)
| 85 | return self.x * other.x + self.y * other.y + self.z * other.z |
| 86 | |
| 87 | def norm(self, epsilon: float = 1e-6) -> Float: |
| 88 | """Compute Norm of Vec3Array, clipped to epsilon.""" |
| 89 | # To avoid NaN on the backward pass, we must use maximum before the sqrt |
| 90 | norm2 = self.dot(self) |
| 91 | if epsilon: |
| 92 | norm2 = jnp.maximum(norm2, epsilon**2) |
| 93 | return jnp.sqrt(norm2) |
| 94 | |
| 95 | def norm2(self): |
| 96 | return self.dot(self) |