How to compute Atan2 for tensors?

[cryscan]

I am trying to implement DeepPhase in candle but I am struggling figuring out how to calculate the phase angles from two tensors using atan2 operation.

[cryscan]

Okay I found out that we can implement custom ops. So is this correct?

use candle_core::{CpuStorage, CustomOp2, Layout, Result, Shape, Tensor};

pub struct Atan2;

impl CustomOp2 for Atan2 {
    fn name(&self) -> &'static str {
        "atan2"
    }

    fn cpu_fwd(
        &self,
        s1: &CpuStorage,
        l1: &Layout,
        s2: &CpuStorage,
        l2: &Layout,
    ) -> Result<(CpuStorage, Shape)> {
        if l1.shape() != l2.shape() {
            candle_core::bail!("operands must have the same shape");
        }

        let s1 = match l1.contiguous_offsets() {
            None => candle_core::bail!("input has to be contiguous"),
            Some((o1, o2)) => &s1.as_slice::<f32>()?[o1..o2],
        };
        let s2 = match l2.contiguous_offsets() {
            None => candle_core::bail!("input has to be contiguous"),
            Some((o1, o2)) => &s2.as_slice::<f32>()?[o1..o2],
        };

        let dst = itertools::zip_eq(s1, s2)
            .map(|(&y, &x)| y.atan2(x))
            .collect();
        let storage = candle_core::WithDType::to_cpu_storage_owned(dst);
        Ok((storage, l1.shape().clone()))
    }

    fn bwd(
        &self,
        y: &Tensor,
        x: &Tensor,
        _: &Tensor,
        _: &Tensor,
    ) -> Result<(Option<Tensor>, Option<Tensor>)> {
        let d = (x.sqr()? + y.sqr()?)?;
        Ok((Some(x.div(&d)?), Some(y.div(&d)?.neg()?))))
    }
}