text: # Copyright      2021  Xiaomi Corp.       (authors: Fangjun Kuang)
  # See ../../../LICENSE for clarification regarding multiple authors
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  
  import torch
  
  
  class _AbsFunction(torch.autograd.Function):
  
      @staticmethod
      def forward(ctx, sparse_tensor: torch.Tensor) -> torch.Tensor:
          '''Compute 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.
          '''
          assert sparse_tensor.is_sparse
          assert sparse_tensor.is_coalesced()
  
          indices = sparse_tensor.indices().clone()
          values = sparse_tensor.values()
          size = sparse_tensor.size()
  
          values_abs = values.abs()
  
          ans = torch.sparse_coo_tensor(indices=indices,
                                        values=values_abs,
                                        size=size,
                                        dtype=sparse_tensor.dtype,
                                        device=sparse_tensor.device)
  
          ctx.save_for_backward(sparse_tensor)
          return ans
  
      @staticmethod
      def backward(ctx, ans_grad: torch.Tensor) -> torch.Tensor:
          sparse_tensor, = ctx.saved_tensors
  
          indices = sparse_tensor.indices().clone()
          values = sparse_tensor.values()
          size = sparse_tensor.size()
  
          sparse_tensor_grad_values = ans_grad.values() * values.sign()
  
          sparse_tensor_grad = torch.sparse_coo_tensor(
              indices=indices,
              values=sparse_tensor_grad_values,
              size=size,
              dtype=sparse_tensor.dtype,
              device=sparse_tensor.device)
  
          return sparse_tensor_grad
  
  
  def abs(sparse_tensor: torch.Tensor) -> torch.Tensor:
      '''Compute 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.
      '''
      return _AbsFunction.apply(sparse_tensor)
  