EnzymeAD/Reactant.jl

Missing neural network operations

Open

#250 创建于 2024年11月9日

在 GitHub 查看
 (0 评论) (0 反应) (0 负责人)Julia (64 fork)github user discovery
good first issue

仓库指标

Star
 (336 star)
PR 合并指标
 (PR 指标待抓取)

描述

NNlib

  • activations
    • tanh_fast
    • sigmoid
    • sigmoid_fast
  • softmax!
  • logsoftmax!
  • conv / conv! (#258)
  • maxpool / maxpool! (#258)
  • meanpool / meanpool! (#258)
  • batched_mul / batched_mul! (#258)
  • padding
    • pad_constant
    • pad_circular
    • pad_repeat
    • pad_zeros
    • pad_reflect
    • pad_symmetric
  • ∇conv_data / ∇conv_data!
  • ∇conv_filter / ∇conv_filter!
  • gather / gather!
  • scatter / scatter! https://github.com/EnzymeAD/Reactant.jl/pull/1395
  • dot_product_attention
  • pixel_shuffle
  • batchnorm (has a special stablehlo impl) - this needs to wait for a corresponding adjoint on EnzymeJAX end
  • grid_sample / grid_sample!
  • upsampling
    • nearest
    • linear
    • bilinear
    • trilinear
  • fold / fold!
  • unfold / unfold!

Some of these might not even be needed, but we should test the generated IR. Strikethrough denotes we don't need specialized handling for these operations

Feel free to add missing operations to the list

贡献者指南