[BACKEND] Add isfinite backward support in aipu (#20)

Author: zhzhcookie

.github/workflows/aipu-build-and-test.yml

@@ -52,7 +52,7 @@ jobs:
           source ~/env_setup.sh
           export FLAGTREE_BACKEND=aipu
           cd python
-          python3.10 -m pip install . --no-build-isolation -v
+          MAX_JOBS=64 python3.10 -m pip install . --no-build-isolation -v
 
       - name: FlagTree Test on AIPU
         shell: bash

third_party/aipu/backend/aipu_torch_dev.cpp

@@ -5,6 +5,7 @@
 #include <c10/util/ArrayRef.h>
 
 #include <torch/csrc/Device.h>
+#include <torch/csrc/autograd/autograd_not_implemented_fallback.h>
 #include <torch/csrc/jit/serialization/pickler.h>
 #include <torch/extension.h>
 
@@ -328,6 +329,11 @@ TORCH_LIBRARY_IMPL(aten, PrivateUse1, m) {
   m.impl("aten::fill_.Scalar", &fill_scalar_aipu);
 }
 
+// Register the autograd dispatch key for operators that have no dispatches
+TORCH_LIBRARY_IMPL(aten, AutogradPrivateUse1, m) {
+  m.impl("isfinite", torch::autograd::autogradNotImplementedFallback());
+}
+
 namespace aipu {
 
 bool is_available() { return context() != nullptr; }