sparse_api.yaml

- api : add
  args : (Tensor x, Tensor y)
  output : Tensor(out)
  kernel :
    func : add_coo_coo{sparse_coo -> sparse_coo},
           add_csr_csr{sparse_csr -> sparse_csr}
    layout : x
  backward : add_grad

- api : conv3d
  args : (Tensor x, Tensor kernel, int[] paddings, int[] dilations, int[] strides, int groups, bool subm)
  output : Tensor(out), Tensor(rulebook)
  kernel :
    func : sparse_conv3d{sparse_coo, dense -> sparse_coo, dense}
    layout : x
  intermediate : rulebook
  backward : conv3d_grad

- api : coo_to_dense
  args : (Tensor x)
  output : Tensor(out)
  invoke : to_dense_impl(x)
  backward : coo_to_dense_grad

- api : create_sparse_coo_tensor
  args : (Tensor values, Tensor indices, IntArray dense_shape)
  output : Tensor(out)
  kernel :
    func : sparse_coo_tensor{dense, dense -> sparse_coo}
    layout : values
    data_type : values
  backward : create_sparse_coo_tensor_grad

- api : dense_to_coo
  args : (Tensor x, int64_t sparse_dim)
  output : Tensor(out)
  invoke : to_sparse_coo_impl(x, sparse_dim)
  backward : dense_to_coo_grad

- api : divide
  args : (Tensor x, Tensor y)
  output : Tensor(out)
  kernel :
    func : divide_coo_coo{sparse_coo -> sparse_coo},
           divide_csr_csr{sparse_csr -> sparse_csr}
    layout : x
  backward : divide_grad

- api : multiply
  args : (Tensor x, Tensor y)
  output : Tensor(out)
  kernel :
    func : multiply_coo_coo{sparse_coo -> sparse_coo},
           multiply_csr_csr{sparse_csr -> sparse_csr}
    layout : x
  backward : multiply_grad

- api : relu
  args : (Tensor x)
  output : Tensor(out)
  kernel :
    func : sparse_coo_relu{sparse_coo -> sparse_coo},
           sparse_csr_relu{sparse_csr -> sparse_csr}
    layout : x
  backward : relu_grad

- api : sin
  args : (Tensor x)
  output : Tensor(out@SparseCooTensor)
  kernel :
    func : sparse_coo_sin {sparse_coo -> sparse_coo},
           sparse_csr_sin {sparse_csr -> sparse_csr}
    layout : x
  backward : sin_grad

- api : softmax
  args : (Tensor x, int axis=-1)
  output : Tensor(out)
  kernel :
    func : softmax_csr{sparse_csr -> sparse_csr}
    layout : x
  backward : softmax_grad

- api : sqrt
  args : (Tensor x)
  output : Tensor(out)
  kernel :
    func : sparse_coo_sqrt{sparse_coo -> sparse_coo},
           sparse_csr_sqrt{sparse_csr -> sparse_csr}
    layout : x
  backward : sqrt_grad

- api : subtract
  args : (Tensor x, Tensor y)
  output : Tensor(out)
  kernel :
    func : subtract_coo_coo{sparse_coo -> sparse_coo},
           subtract_csr_csr{sparse_csr -> sparse_csr}
    layout : x
  backward : subtract_grad

- api : tanh
  args : (Tensor x)
  output : Tensor(out)
  kernel :
    func : sparse_coo_tanh{sparse_coo -> sparse_coo},
           sparse_csr_tanh{sparse_csr -> sparse_csr}
    layout : x
  backward : tanh_grad

- api : to_dense
  args : (Tensor x)
  output : Tensor(out)
  invoke : to_dense_impl(x)

- api : to_sparse_coo
  args : (Tensor x, int64_t sparse_dim)
  output : Tensor(out)
  invoke : to_sparse_coo_impl(x, sparse_dim)

- api : to_sparse_csr
  args : (Tensor x)
  output : Tensor(out)
  invoke : to_sparse_csr_impl(x)

- api : values
  args : (Tensor x)
  output : Tensor(out)
  kernel :
    func : coo_values{sparse_coo -> dense},
           csr_values{sparse_csr -> dense}
    layout : x
  backward : values_grad

- api: full_like
  args : (Tensor x, Scalar value, DataType dtype=DataType::UNDEFINED)
  output : Tensor(out)
  kernel :
    func : coo_full_like{sparse_coo -> sparse_coo},
           csr_full_like{sparse_csr -> sparse_csr}
    layout : x
    data_type : dtype

- api: fused_attention
  args : (Tensor query, Tensor key, Tensor value, Tensor sparse_mask, Tensor key_padding_mask, Tensor attn_mask)
  output : Tensor(out), Tensor(softmax)
  kernel :
    func : fused_attention_csr{dense, dense, dense, sparse_csr, dense, dense -> dense, sparse_csr}
    layout : sparse_mask
  intermediate : softmax
  backward: fused_attention_grad

- api: masked_matmul
  args : (Tensor x, Tensor y, Tensor mask)
  output : Tensor(out)
  kernel :
    func : csr_masked_matmul{dense, dense, sparse_csr -> sparse_csr}
    layout : x
  backward: masked_matmul_grad

- api: matmul
  args : (Tensor x, Tensor y)
  output : Tensor(out)
  kernel :
    func : csr_dense_matmul{sparse_csr, dense -> dense},
           csr_csr_matmul{sparse_csr, sparse_csr -> sparse_csr},
           coo_dense_matmul{sparse_coo, dense -> dense},
           coo_coo_matmul{sparse_coo, sparse_coo -> sparse_coo}
    layout : x
  backward: matmul_grad

- api: maxpool
  args : (Tensor x, int[] kernel_sizes, int[] paddings, int[] dilations, int[] strides)
  output : Tensor(out), Tensor(rulebook)
  kernel :
    func : sparse_maxpool{sparse_coo -> sparse_coo, dense}
    layout : x
  intermediate : rulebook
  backward : sparse_maxpool_grad

- api: mv
  args : (Tensor x, Tensor vec)
  output : Tensor(out)
  kernel :
    func : mv_coo{sparse_coo, dense -> dense},
           mv_csr{sparse_csr, dense -> dense}
    layout : x
  backward: mv_grad