o Hi @s6ddlZGdddejjZdejdejfddZdS)Nc@s@eZdZedejdejfddZedejdejfddZdS) _AbsFunction sparse_tensorreturncCs`|jsJ|s J|}|}|}|}tj||||j |j d}| ||S)a8Compute the `abs` of a sparse tensor. Args: sparse_tensor: A sparse tensor. It has to satisfy:: assert sparse_tensor.is_coalesced() Returns: The absolute value of the sparse tensor. The `abs` operation is applied element-wise. indicesvaluessizedtypedevice) is_sparse is_coalescedrclonerrabstorchsparse_coo_tensorr r save_for_backward)ctxrrrr values_absansrF/home/ubuntu/.local/lib/python3.10/site-packages/k2/sparse/autograd.pyforwards    z_AbsFunction.forwardans_gradcCsP|j\}|}|}|}||}tj||||j|j d}|S)Nr) saved_tensorsrr rrsignrrr r )rrrrrrsparse_tensor_grad_valuessparse_tensor_gradrrrbackward5s z_AbsFunction.backwardN)__name__ __module__ __qualname__ staticmethodrTensorrrrrrrrs rrrcCs t|S)a3Compute the `abs` of a sparse tensor. It supports autograd. Args: sparse_tensor: A sparse tensor. It has to satisfy:: assert sparse_tensor.is_coalesced() Returns: The absolute value of the sparse tensor. The `abs` operation is applied element-wise. )rapply)rrrrrIs r)rautogradFunctionrr"rrrrrs6